Files
wehub-resource-sync 7254f7b4d1
Build / Build (macos-latest) (push) Has been cancelled
Build / Build (ubuntu-latest) (push) Has been cancelled
Build / Build (windows-latest) (push) Has been cancelled
Build / Analyze (javascript) (push) Has been cancelled
Build / Analyze (python) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:37:45 +08:00

28386 lines
1.1 MiB
Plaintext

// Experimental
import * as base from './base.js';
import * as text from './text.js';
const mlir = {};
const _ = {};
mlir.ModelFactory = class {
async match(context) {
const stream = context.stream;
const identifier = context.identifier;
const extension = identifier.split('.').pop().toLowerCase();
if (stream && stream.length > 4) {
const buffer = stream.peek(4);
const signature = String.fromCharCode.apply(null, buffer);
if (signature === 'ML\xEFR') {
return context.set('mlir.binary');
}
}
try {
const reader = await context.read('text', 0x10000);
let whitespace = true;
for (let line = reader.read('\n'); line !== undefined; line = reader.read('\n')) {
if (/module\s+(@\w+|\w+|attributes|\{)/.test(line) ||
/tensor<[\w\d]+>/.test(line) ||
/func[.\s]*@\w+/.test(line) ||
/%\w+\s*=\s*"[\w.]+/.test(line) ||
/%\w+\s*=\s*\w+\./.test(line) ||
/!\w+\s*=\s*![\w.]+</.test(line) ||
/#\w+\s*=\s*#[\w.]+</.test(line) ||
/#\w+\s*=\s*loc\s*\(/.test(line) ||
/\w+\.\w+(?:\s+\w+)*\s+@\w+/.test(line) ||
/\w+\.\w+\s+#[\w.]+</.test(line) ||
/\w+\.\w+\s*<?\{/.test(line) ||
/:\s*![\w.]+/.test(line) ||
/(%\w+|\w{2,}|[)])\s*:\s*(\[|tensor<)/.test(line) ||
/->\s*(![\w.]+|\(|tensor<)/.test(line)) {
return context.set('mlir.text');
}
if (line && !line.trim().startsWith('//')) {
whitespace = false;
}
}
if (extension === 'mlir' && whitespace) {
return context.set('mlir.text');
}
} catch {
// continue
}
return null;
}
async open(context) {
const metadata = await mlir.Metadata.open(context);
switch (context.type) {
case 'mlir.text': {
const decoder = await context.read('text.decoder');
const config = new _.ParserConfig(new _.DialectContext(metadata));
const sourceMgr = new _.SourceMgr();
const buffer = new _.MemoryBuffer(decoder, context.identifier);
sourceMgr.addNewSourceBuffer(buffer);
const state = new _.ParserState(sourceMgr, config);
const parser = new _.TopLevelOperationParser(state);
const block = new _.Block();
parser.parse(block);
const model = new mlir.Model(config, 'MLIR', '', block, state.attributeAliasDefinitions);
return model;
}
case 'mlir.binary': {
const buffer = await context.read('binary');
const config = new _.ParserConfig(new _.DialectContext(metadata));
const reader = new _.BytecodeReader(buffer, config);
const block = reader.read();
const format = `MLIR Bytecode v${reader.version.value}`;
const producer = reader.producer;
const model = new mlir.Model(config, format, producer, block, new Map());
return model;
}
default: {
throw new mlir.Error(`Unsupported MLIR format '${context.type}'.`);
}
}
}
};
mlir.Model = class {
constructor(config, format, producer, block, attributeAliasDefinitions) {
this.format = format;
this.producer = producer || '';
this.modules = [];
this.functions = [];
this.metadata = [];
const modules = [];
const isFunc = (name) => name.endsWith('.func') || /\.func_v\d+$/.test(name);
const isModule = (name) => name.endsWith('.module');
const collectModules = (operations, path, attributes) => {
let identifier = 0;
const funcs = [];
const ops = [];
for (const op of operations) {
if (isFunc(op.name.getStringRef())) {
funcs.push(op);
} else if (isModule(op.name.getStringRef())) {
let name = op.getAttr('sym_name');
name = name ? name.value : `$${identifier++}`;
const modulePath = [...path, name];
for (const region of op.regions || []) {
for (const blk of region.blocks || []) {
collectModules(blk.operations || [], modulePath, op.getAttrDictionary());
}
}
} else {
ops.push(op);
}
}
if (funcs.length > 0 || ops.length > 0) {
let name = null;
if (attributes.get('sym_name')) {
name = attributes.get('sym_name');
name = `@${name.value}`;
}
modules.push({ path, symName: name, funcs, ops, attributes });
}
};
collectModules(block.operations, [], new Map());
const formatPrefix = (path, symName) => {
if (symName) {
return symName;
}
if (modules.length !== 1 && path.length > 0) {
return path.map((path) => `${path}`).join('::');
}
return '';
};
const functions = new Map();
let identifier = 0;
for (const module of modules) {
const prefix = formatPrefix(module.path, module.symName);
for (const func of module.funcs) {
const sym_name = func.getAttr('sym_name');
const base = sym_name ? sym_name.value : `$${identifier}`;
identifier++;
const name = prefix ? `${prefix}::@${base}` : `@${base}`;
functions.set(name, { func, prefix, base, module });
}
}
const context = new mlir.Context(functions);
for (const [name, info] of functions) {
const graph = context.graph(info.func, name);
this.functions.push(graph);
}
for (const module of modules) {
if (module.ops.length > 0 || module.attributes.size > 0) {
const name = formatPrefix(module.path, module.symName) || '';
const opName = _.RegisteredOperationName.lookup('builtin.module', config.context);
const state = new _.OperationState(null, opName);
state.attributes = module.attributes;
state.regions = [{ blocks: [{ operations: module.ops, arguments: [] }] }];
const op = _.Operation.create(state);
const graph = context.graph(op, name);
this.modules.push(graph);
}
}
for (const [name, attribute] of attributeAliasDefinitions) {
let value = attribute.type;
if (!value) {
value = attribute.toString();
}
const metadata = new mlir.Argument(name, value, 'attribute');
this.metadata.push(metadata);
}
}
};
mlir.Graph = class {
constructor(func, context, name) {
this.name = name || '';
if (!name && func.attributes.has('sym_name')) {
const sym_name = func.attributes.get('sym_name');
this.name = sym_name.value;
}
this.type = 'graph';
if (func.name === 'func' || func.name.getStringRef().endsWith('.func') || /\.func_v\d+$/.test(func.name.getStringRef())) {
this.type = 'function';
}
this.inputs = [];
this.outputs = [];
this.nodes = [];
this.metadata = [];
const tensors = new Map();
if (func.attributes.has('function_type')) {
const function_type = func.attributes.get('function_type');
const args = func.regions && func.regions[0] && func.regions[0].blocks && func.regions[0].blocks[0] && func.regions[0].blocks[0].arguments ? func.regions[0].blocks[0].arguments : [];
const inputs = function_type.type.inputs;
const results = function_type.type.results;
for (let i = 0; i < inputs.length; i++) {
const input = inputs[i];
const name = args[i] && args[i].name ? args[i].name : `%arg${i}`;
const type = mlir.Utility.valueType(input.type || input);
const value = new mlir.Value(name, type, '', null);
const argument = new mlir.Argument(name, [value]);
this.inputs.push(argument);
}
for (let i = 0; i < results.length; i++) {
const output = results[i];
const name = output.value || i.toString();
const type = mlir.Utility.valueType(output.type);
const valueName = output.value || output.name || `%result${i}`;
const value = new mlir.Value(valueName, type, '', null);
const argument = new mlir.Argument(name, [value]);
this.outputs.push(argument);
}
}
const values = new Map();
values.map = (name) => {
if (!values.has(name)) {
values.set(name, { name, to: [], from: [] });
}
return values.get(name);
};
const operations = [];
for (const region of func.regions) {
for (const block of region.blocks) {
for (const op of block.operations) {
const operation = {
identifier: op.identifier,
name: op.name,
label: op.label,
attributes: op.getAttrDictionary(),
operands: [],
results: [],
inputs: [],
outputs: [],
regions: op.regions,
loc: op.loc,
delete: false,
};
const opMetadata = op.name.metadata;
const operands = op.operands;
let lastVariadicIndex = -1;
let lastVariadicName = null;
if (opMetadata && opMetadata.operands) {
for (let j = opMetadata.operands.length - 1; j >= 0; j--) {
const metaOp = opMetadata.operands[j];
if (metaOp.type && metaOp.type.name === 'Variadic') {
lastVariadicIndex = j;
lastVariadicName = metaOp.name;
break;
}
}
}
for (let i = 0; i < operands.length; i++) {
const input = op.operands[i];
let inputName = null;
const isVariadicOverflow = lastVariadicIndex >= 0 && i >= lastVariadicIndex;
if (opMetadata && opMetadata.operands && opMetadata.operands[i]) {
inputName = opMetadata.operands[i].name;
} else if (isVariadicOverflow) {
inputName = lastVariadicName;
} else {
inputName = input.name || i.toString();
}
if (typeof input.name !== 'string' || !input.name) {
throw new mlir.Error(`Invalid operand name '${JSON.stringify(input.name)}'.`);
}
const value = values.map(input.name);
value.to.push(operation);
const arg = { name: input.name, type: input.type };
operation.operands.push(arg);
if (isVariadicOverflow && operation.inputs.length > 0 && operation.inputs[operation.inputs.length - 1].name === inputName) {
operation.inputs[operation.inputs.length - 1].value.push(arg);
} else {
operation.inputs.push({ name: inputName, value: [arg] });
}
}
const results = op.results;
let lastVariadicResultIndex = -1;
let lastVariadicResultName = null;
if (opMetadata && opMetadata.results) {
for (let j = opMetadata.results.length - 1; j >= 0; j--) {
const metaRes = opMetadata.results[j];
if (metaRes.type && metaRes.type.name === 'Variadic') {
lastVariadicResultIndex = j;
lastVariadicResultName = metaRes.name;
break;
}
}
}
for (let i = 0; i < results.length; i++) {
const output = results[i];
if (!output.name) {
continue;
}
const value = values.map(output.name);
value.type = output.type;
value.from.push(operation);
let outputName = null;
const isVariadicOverflow = lastVariadicResultIndex >= 0 && i >= lastVariadicResultIndex;
if (opMetadata && opMetadata.results && opMetadata.results[i]) {
outputName = opMetadata.results[i].name;
} else if (isVariadicOverflow) {
outputName = lastVariadicResultName;
} else {
outputName = output.name;
}
operation.results.push(value);
if (isVariadicOverflow && operation.outputs.length > 0 && operation.outputs[operation.outputs.length - 1].name === outputName) {
operation.outputs[operation.outputs.length - 1].value.push(value);
} else {
operation.outputs.push({
name: outputName,
value: [value]
});
}
}
operations.push(operation);
}
}
}
const constantMap = new Map();
const constantTypes = new Set([
'tosa.const', 'stablehlo.constant', 'arith.constant',
'mhlo.constant', 'torch.constant.tensor', 'onnx.Constant',
'vhlo.constant_v1'
]);
for (const op of operations) {
if (constantTypes.has(op.name.getStringRef()) &&
op.operands.length === 0 &&
op.attributes.size === 1 &&
op.results.length === 1) {
const result = op.results[0];
if (result.to && result.to.length === 1) {
if (result.to[0].name.getStringRef().endsWith('.return') || result.to[0].name.getStringRef() === 'vhlo.return_v1') {
continue;
}
const valueAttr = op.attributes.get('value') || op.attributes.get('values');
if ((valueAttr instanceof _.DenseElementsAttr || valueAttr instanceof _.DenseResourceElementsAttr) && valueAttr.value !== null && valueAttr.type && valueAttr.type instanceof _.RankedTensorType) {
const type = mlir.Utility.valueType(valueAttr.type);
if (type instanceof mlir.TensorType && !type.dataType.startsWith('!')) {
const data = valueAttr.value;
constantMap.set(result.name, new mlir.Tensor(type, data));
op.delete = true;
}
}
}
}
}
const torchConstantMap = new Map();
for (const op of operations) {
const opName = op.name.getStringRef();
if (opName === 'torch.constant.int' ||
opName === 'torch.constant.bool' ||
opName === 'torch.constant.float' ||
opName === 'torch.constant.str' ||
opName === 'torch.constant.none') {
if (op.operands.length === 0 &&
op.results.length === 1) {
const result = op.results[0];
if (result.to && result.to.length === 1) {
let value = null;
let type = null;
const attr = op.attributes.get('value');
const attrValue = attr && typeof attr === 'object' ? attr.value : attr;
if (opName === 'torch.constant.int') {
value = attrValue === undefined ? 0 : attrValue;
type = 'int64';
} else if (opName === 'torch.constant.bool') {
value = attrValue === undefined ? false : attrValue;
type = 'boolean';
} else if (opName === 'torch.constant.float') {
value = attrValue === undefined ? 0.0 : attrValue;
type = 'float64';
} else if (opName === 'torch.constant.str') {
value = attrValue === undefined ? '' : attrValue;
type = 'string';
} else if (opName === 'torch.constant.none') {
value = null;
type = 'none';
}
torchConstantMap.set(result.name, { value, type });
op.delete = true;
}
}
}
}
for (const op of operations) {
const opName = op.name.getStringRef();
if (opName === 'torch.prim.ListConstruct' && op.results.length === 1) {
const result = op.results[0];
if (result.to && result.to.length === 1) {
const inputValues = [];
let allConstant = true;
for (const operand of op.operands) {
if (torchConstantMap.has(operand.name)) {
inputValues.push(torchConstantMap.get(operand.name).value);
} else {
allConstant = false;
break;
}
}
if (allConstant) {
torchConstantMap.set(result.name, { value: inputValues, type: 'list' });
op.delete = true;
}
}
}
}
const tensor = (arg) => {
if (!tensors.has(arg.name)) {
const initializer = constantMap.get(arg.name) || null;
const type = arg.type ? mlir.Utility.valueType(arg.type) : null;
tensors.set(arg.name, new mlir.Value(arg.name, type, null, initializer));
}
return tensors.get(arg.name);
};
for (const input of this.inputs) {
for (const arg of input.value) {
if (!tensors.has(arg.name)) {
tensors.set(arg.name, arg);
}
}
}
const returnOp = operations.find((op) => op.name.getStringRef().endsWith('.return') || op.name.getStringRef() === 'vhlo.return_v1');
if (returnOp) {
for (let i = 0; i < this.outputs.length && i < returnOp.operands.length; i++) {
const returnValue = returnOp.operands[i];
if (returnValue && typeof returnValue.name === 'string' && returnValue.name.startsWith('%')) {
const output = this.outputs[i];
const returnType = mlir.Utility.valueType(returnValue.type);
const initializer = constantMap.get(returnValue.name) || null;
output.value[0] = new mlir.Value(returnValue.name, returnType, '', initializer);
}
}
returnOp.delete = true;
}
for (const output of this.outputs) {
for (let i = 0; i < output.value.length; i++) {
const arg = output.value[i];
if (tensors.has(arg.name)) {
output.value[i] = tensors.get(arg.name);
} else {
tensors.set(arg.name, arg);
}
}
}
for (const op of operations.filter((op) => !op.delete)) {
const node = new mlir.Node(op, context, tensor, torchConstantMap);
this.nodes.push(node);
}
for (const [name, value] of func.attributes) {
if (name === 'sym_name' || name === 'function_type') {
continue;
}
const v = value.value === undefined ? value.toString() : value.value;
const metadata = new mlir.Argument(name, v, value.type ? mlir.Utility.argumentType(value.type) : 'attribute');
this.metadata.push(metadata);
}
}
};
mlir.Argument = class {
constructor(name, value, type = null) {
this.name = name;
this.value = value;
this.type = type;
}
};
mlir.Value = class {
constructor(name, type, description, initializer) {
if (typeof name !== 'string') {
throw new mlir.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
}
this.name = name;
this.type = !type && initializer ? initializer.type : type;
this.description = description || null;
this.initializer = initializer || null;
}
};
mlir.Node = class {
constructor(op, context, tensor, torchConstantMap) {
if (!op.name) {
throw new mlir.Error('Undefined node type.');
}
this.name = '';
this.type = { ...op.name.getRegisteredInfo()?.metadata };
this.type.name = op.label || op.identifier || '';
this.type.identifier = op.name.getStringRef() || '';
this.inputs = [];
this.outputs = [];
this.attributes = [];
this.blocks = [];
this.metadata = [];
if (op.loc instanceof _.LocationAttr && op.loc instanceof _.UnknownLoc === false) {
this.metadata.push(new mlir.Argument('location', op.loc.printLocationInternal(), 'attribute'));
}
torchConstantMap = torchConstantMap || new Map();
const segmentSizes = op.attributes.get('operandSegmentSizes');
const operandMeta = this.type && this.type.operands;
const operands = op.inputs || [];
if (segmentSizes && operandMeta && segmentSizes.value.length === operandMeta.length) {
let offset = 0;
for (let i = 0; i < segmentSizes.value.length; i++) {
const size = segmentSizes.value[i];
const name = operandMeta[i].name;
for (let j = 0; j < size; j++) {
if (offset + j < operands.length) {
operands[offset + j] = { ...operands[offset + j], name };
}
}
offset += size;
}
}
const operandGroups = new Map();
const operandOrder = [];
for (const input of operands) {
if (!operandGroups.has(input.name)) {
operandGroups.set(input.name, []);
operandOrder.push(input.name);
}
operandGroups.get(input.name).push(input);
}
for (const name of operandOrder) {
const inputs = operandGroups.get(name);
let argument = null;
if (inputs.length === 1) {
const [input] = inputs;
if (Array.isArray(input.value) && input.value.length === 1) {
const val = input.value[0];
if (val && typeof val.name === 'string' && torchConstantMap.has(val.name)) {
const constant = torchConstantMap.get(val.name);
argument = new mlir.Argument(input.name, constant.value, constant.type);
this.inputs.push(argument);
continue;
}
}
if (input.type) {
const argumentType = mlir.Utility.argumentType(input.type);
if (argumentType === 'tensor') {
const tensorType = mlir.Utility.valueType(input.type);
const value = new mlir.Tensor(tensorType, input.value);
argument = new mlir.Argument(input.name, value, 'tensor');
} else {
argument = new mlir.Argument(input.name, input.value, argumentType);
}
} else if (Array.isArray(input.value) && !input.value.every((value) => typeof value.name === 'string' && value.name.startsWith('%'))) {
argument = new mlir.Argument(input.name, input.value, input.type ? mlir.Utility.argumentType(input.type) : 'attribute');
} else if (Array.isArray(input.value)) {
argument = new mlir.Argument(input.name, input.value.map((arg) => tensor(arg)));
} else {
argument = new mlir.Argument(input.name, input.value, input.type ? mlir.Utility.argumentType(input.type) : 'attribute');
}
} else {
let allConstants = true;
const constantValues = [];
for (const input of inputs) {
if (Array.isArray(input.value)) {
for (const arg of input.value) {
if (arg && typeof arg.name === 'string' && torchConstantMap.has(arg.name)) {
constantValues.push(torchConstantMap.get(arg.name).value);
} else {
allConstants = false;
break;
}
}
} else {
allConstants = false;
}
if (!allConstants) {
break;
}
}
if (allConstants && constantValues.length > 0) {
argument = new mlir.Argument(name, constantValues, 'list');
} else {
const values = [];
for (const input of inputs) {
if (Array.isArray(input.value)) {
values.push(...input.value.map((arg) => tensor(arg)));
} else {
values.push(tensor({ name: input.value, type: input.type }));
}
}
argument = new mlir.Argument(name, values);
}
}
this.inputs.push(argument);
}
for (const output of op.outputs || []) {
const argument = new mlir.Argument(output.name, output.value.map((arg) => tensor(arg)));
this.outputs.push(argument);
}
if (op.attributes) {
for (const [name, attr] of op.attributes) {
let value = attr;
let type = null;
if (attr instanceof _.SymbolRefAttr && context) {
const graph = context.function(`${value.value}`);
if (graph) {
value = graph;
type = 'function';
}
} else if (name === 'callee' && attr instanceof _.StringAttr && context) {
const graph = context.function(`@${attr.value}`) || context.function(attr.value);
if (graph) {
value = graph;
type = 'function';
}
} else if (attr instanceof _.DenseElementsAttr) {
const tensorType = mlir.Utility.valueType(attr.type);
if (tensorType instanceof mlir.TensorType && tensorType.dataType.startsWith('!')) {
value = `dense<${attr.type}>`;
} else {
value = new mlir.Tensor(tensorType, attr.value);
type = 'tensor';
}
} else if (attr instanceof _.DenseResourceElementsAttr) {
value = new mlir.Tensor(mlir.Utility.valueType(attr.type), attr.value);
type = 'tensor';
} else if (attr instanceof _.SparseElementsAttr) {
value = new mlir.Tensor(mlir.Utility.valueType(attr.type), attr.values, attr.indices);
type = 'tensor';
} else if (attr instanceof _.DenseArrayAttr) {
value = attr.value;
} else if (Array.isArray(attr)) {
value = attr;
} else if (attr) {
value = attr.toString();
}
const attribute = new mlir.Argument(name, value, type || 'attribute');
this.attributes.push(attribute);
}
}
if (op.regions && op.regions.length > 0) {
const opMetadata = this.type;
for (let i = 0; i < op.regions.length; i++) {
const region = op.regions[i];
if (region.blocks && region.blocks.length > 0) {
const name = (opMetadata.regions && opMetadata.regions[i] ? opMetadata.regions[i].name : null) || i.toString();
const blockName = region.blocks[0].name || '';
const func = { name: 'func', attributes: new Map(), regions: [region] };
const graph = new mlir.Graph(func, context, blockName);
const argument = new mlir.Argument(name, graph, 'graph');
this.blocks.push(argument);
}
}
}
}
};
mlir.Tensor = class {
constructor(type, values, indices) {
if (indices !== undefined && type instanceof mlir.TensorType && type.shape && Array.isArray(type.shape.dimensions) && type.shape.dimensions.every((d) => (typeof d === 'bigint' && d >= 0n) || (typeof d === 'number' && d >= 0))) {
const rank = type.shape.dimensions.length;
const numPoints = values ? values.length : 0;
const valuesShape = new mlir.TensorShape([numPoints]);
const indicesShape = rank > 0 && indices.length === numPoints * rank ? new mlir.TensorShape([numPoints, rank]) : new mlir.TensorShape([indices.length]);
this.type = new mlir.TensorType(type.elementType, type.shape, 'sparse');
this.values = new mlir.Tensor(new mlir.TensorType(type.elementType, valuesShape), values || []);
this.indices = new mlir.Tensor(new mlir.TensorType(new _.IntegerType('i64'), indicesShape), indices);
this.encoding = '|';
} else {
this.type = type;
this.values = values;
this.encoding = values instanceof Uint8Array ? '<' : '|';
if (this.encoding === '<' && type.shape && type.shape.dimensions.length > 0) {
const dims = type.shape.dimensions;
const numElements = dims.reduce((a, b) => a * (typeof b === 'bigint' ? Number(b) : b), 1);
const elementType = type.elementType;
let bitWidth = 8;
if (elementType instanceof _.ComplexType && typeof elementType.elementType.getIntOrFloatBitWidth === 'function') {
bitWidth = elementType.elementType.getIntOrFloatBitWidth() * 2;
} else if (elementType && typeof elementType.getIntOrFloatBitWidth === 'function') {
bitWidth = elementType.getIntOrFloatBitWidth();
}
const itemsize = Math.ceil(bitWidth / 8);
if (numElements > 1 && values.length < numElements * itemsize) {
this.stride = new Array(dims.length).fill(0);
}
}
}
}
};
mlir.TensorType = class {
constructor(elementType, shape, layout) {
this.elementType = elementType || null;
this.dataType = elementType ? mlir.Utility.dataType(elementType) : null;
this.shape = shape || new mlir.TensorShape([]);
this.layout = layout || null;
}
toString() {
return this.dataType + this.shape.toString();
}
};
mlir.TensorShape = class {
constructor(dimensions) {
this.dimensions = dimensions;
}
toString() {
if (!this.dimensions || this.dimensions.length === 0) {
return '';
}
return `[${this.dimensions.map((dimension) => dimension.toString()).join(', ')}]`;
}
};
mlir.Context = class {
constructor(functions) {
this._functions = functions; // Map of fullName -> {func, prefix, base, module}
this._graphs = new Map();
this._constructing = new Set();
}
graph(module, name) {
if (!this._graphs.has(name)) {
this._constructing.add(name);
const graph = new mlir.Graph(module, this, name);
this._graphs.set(name, graph);
this._constructing.delete(name);
}
return this._graphs.get(name);
}
function(name) {
if (this._graphs.has(name)) {
return this._graphs.get(name);
}
if (this._constructing.has(name)) {
return { name, type: 'function', nodes: [], inputs: [], outputs: [] };
}
if (this._functions.has(name)) {
const info = this._functions.get(name);
return this.graph(info.func, name);
}
for (const [fullName, info] of this._functions) {
if (info.base === name) {
if (this._graphs.has(fullName)) {
return this._graphs.get(fullName);
}
if (this._constructing.has(fullName)) {
return { name: fullName, type: 'function', nodes: [], inputs: [], outputs: [] };
}
return this.graph(info.func, fullName);
}
}
return null;
}
};
mlir.Utility = class {
static argumentType(type) {
if (type instanceof _.RankedTensorType || type instanceof _.UnrankedTensorType || type instanceof _.torch.ValueTensorType) {
return 'tensor';
}
if (type instanceof _.FunctionType) {
return 'function';
}
if (type instanceof _.NoneType) {
return 'none';
}
if (type instanceof _.IntegerType || type instanceof _.FloatType || type instanceof _.IndexType || type instanceof _.ComplexType) {
return mlir.Utility.dataType(type);
}
if (type instanceof _.VectorType || type instanceof _.MemRefType || type instanceof _.UnrankedMemRefType) {
return type.toString();
}
if (type instanceof _.TupleType) {
return type.toString();
}
if (type instanceof _.Type) {
return type.toString();
}
throw new mlir.Error(`Unexpected argument type '${type}'.`);
}
static dataType(value) {
if (value instanceof _.IntegerType) {
const name = value.toString();
const match = name.match(/^(s|u)?i([0-9]+)$/);
if (match) {
const width = parseInt(match[2], 10);
if (match[1] === 'u') {
return `uint${width}`;
}
return `int${width}`;
}
throw new mlir.Error(`Unexpected integer type '${name}'.`);
}
if (value instanceof _.FloatType) {
const name = value.toString();
switch (name) {
case 'f16': return 'float16';
case 'f32': return 'float32';
case 'f64': return 'float64';
case 'f80': return 'float80';
case 'f128': return 'float128';
case 'bf16': return 'bfloat16';
case 'tf32': return 'tf32';
case 'float8': return 'float8';
case 'fp8': return 'float8';
case 'fp8e4m3': return 'float8e4m3';
case 'fp8_e4m3': return 'float8e4m3';
case 'fp8e4m3fn': return 'float8e4m3fn';
case 'fp8e5m2': return 'float8e5m2';
case 'fp8_e5m2': return 'float8e5m2';
case 'f4E2M1FN': return 'float4e2m1fn';
case 'f6E2M3FN': return 'float6e2m3fn';
case 'f6E3M2FN': return 'float6e3m2fn';
case 'f8E3M4': return 'float8e3m4';
case 'f8E4M3': return 'float8e4m3';
case 'f8E4M3B11FNUZ': return 'float8e4m3b11fnuz';
case 'f8E4M3FN': return 'float8e4m3fn';
case 'f8E4M3FNUZ': return 'float8e4m3fnuz';
case 'f8E5M2': return 'float8e5m2';
case 'f8E5M2FNUZ': return 'float8e5m2fnuz';
case 'f8E8M0FNU': return 'float8e8m0fnu';
default: throw new mlir.Error(`Unexpected float type '${name}'.`);
}
}
if (value instanceof _.IndexType) {
return 'int64';
}
if (value instanceof _.ComplexType) {
return `complex<${mlir.Utility.dataType(value.elementType)}>`;
}
if (value instanceof _.NoneType) {
return 'none';
}
if (value instanceof _.VectorType || value instanceof _.MemRefType || value instanceof _.UnrankedMemRefType) {
return value.toString();
}
if (value instanceof _.TupleType) {
return value.toString();
}
if (value instanceof _.Type) {
return value.toString();
}
throw new mlir.Error(`Unexpected data type '${value}'.`);
}
static valueType(type) {
if (type === undefined) {
return null;
}
if (type instanceof _.RankedTensorType) {
const shape = type.shape.map((d) => d < 0n ? '?' : d);
return new mlir.TensorType(type.elementType, new mlir.TensorShape(shape));
}
if (type instanceof _.VectorType) {
const shape = type.shape.map((d) => d < 0n ? '?' : d);
return new mlir.TensorType(type.elementType, new mlir.TensorShape(shape));
}
if (type instanceof _.MemRefType) {
const shape = type.shape.map((d) => d < 0n ? '?' : d);
return new mlir.TensorType(type.elementType, new mlir.TensorShape(shape));
}
if (type instanceof _.UnrankedTensorType || type instanceof _.UnrankedMemRefType) {
return new mlir.TensorType(type.elementType, null);
}
if (type instanceof _.torch.ValueTensorType) {
const shape = type.optionalSizes ? type.optionalSizes.map((s) => s === -1 ? '?' : BigInt(s)) : null;
return new mlir.TensorType(type.optionalDtype, shape ? new mlir.TensorShape(shape) : null);
}
if (type instanceof _.TupleType) {
return type.toString();
}
if (type instanceof _.Type) {
return type.toString();
}
throw new mlir.Error(`Unexpected value type '${type}'.`);
}
};
_.Block = class {
constructor() {
this.operations = [];
}
};
_.OperationState = class {
constructor(location, name) {
this.location = location;
this.name = name;
if (this.name instanceof _.OperationName === false) {
throw new mlir.Error('Invalid operation name.');
}
this.identifier = name.identifier || name.getStringRef();
delete name.identifier;
this.attributes = new Map();
this.operands = [];
this.types = [];
this.regions = [];
this.propertiesAttr = null;
}
addRegion() {
const region = {};
this.regions.push(region);
return region;
}
addTypes(newTypes) {
if (!Array.isArray(newTypes) || !newTypes.every((type) => type instanceof _.Type)) {
throw new mlir.Error(`Invalid types '${JSON.stringify(newTypes.filter((type) => type instanceof _.Type === false))}'.`);
}
for (const type of newTypes) {
this.types.push(type);
}
}
addAttribute(name, value) {
if (typeof name !== 'string' || name.length === 0) {
throw new mlir.Error(`Invalid attribute name '${JSON.stringify(name)}'.`);
}
this.attributes.set(name, value);
}
getAttr(name) {
if (this.propertiesAttr instanceof _.DictionaryAttr) {
const value = this.propertiesAttr.get(name);
if (value !== undefined) {
return value;
}
}
return this.attributes.get(name);
}
getAttrDictionary() {
if (this.propertiesAttr instanceof _.DictionaryAttr) {
const result = new Map(this.attributes);
for (const [name, value] of this.propertiesAttr.value) {
result.set(name, value);
}
return result;
}
return this.attributes;
}
};
_.OperationName = class {
constructor(dialect, name) {
this.dialect = dialect;
this.name = name;
}
getStringRef() {
return this.name;
}
getRegisteredInfo() {
return null;
}
};
_.RegisteredOperationName = class extends _.OperationName {
constructor(dialect, name, metadata) {
super(dialect, name);
this.metadata = metadata;
if (metadata.assemblyFormat) {
const parser = new _.AssemblyFormatParser(metadata);
this.directives = parser.parse();
}
}
static lookup(name, context) {
const index = name.indexOf('.');
if (index !== -1) {
const dialectName = name.substring(0, index);
const dialect = context.getOrLoadDialect(dialectName);
if (dialect) {
return dialect.getOperation(name);
}
}
return null;
}
getRegisteredInfo() {
return this;
}
hasTrait(type) {
if (this.metadata && Array.isArray(this.metadata.traits)) {
for (const trait of this.metadata.traits) {
if (trait.type && trait.type.name === type) {
return true;
}
}
}
return false;
}
};
_.Operation = class {
static create(state) {
return new _.Operation(state);
}
constructor(state) {
this.name = state.name; // registered operation name
this.identifier = state.identifier; // original parsed identifier
this.label = state.label; // human-readable name
this.attributes = state.attributes;
this.operands = state.operands;
this.regions = state.regions;
this.propertiesAttr = state.propertiesAttr;
this.loc = state.location;
this.results = [];
if (Array.isArray(state.types)) {
for (let i = 0; i < state.types.length; i++) {
const result = new _.OpResult(this, i, state.types[i]);
this.results.push(result);
}
}
}
getAttr(name) {
if (this.propertiesAttr instanceof _.DictionaryAttr) {
const value = this.propertiesAttr.get(name);
if (value !== undefined) {
return value;
}
}
return this.attributes.get(name);
}
getAttrDictionary() {
if (this.propertiesAttr instanceof _.DictionaryAttr) {
const result = new Map(this.attributes);
for (const [key, value] of this.propertiesAttr.value) {
result.set(key, value);
}
return result;
}
return this.attributes;
}
};
_.UnresolvedOperand = class {
constructor(location, name, number) {
this.location = location;
this.name = name;
this.number = number;
}
toString() {
return this.number > 0 ? `${this.name}#${this.number}` : this.name;
}
};
_.Value = class {
constructor(name, type) {
this.name = name;
this.type = type;
}
toString() {
return this.name;
}
};
_.OpResult = class extends _.Value {
constructor(owner, resultNumber, type) {
super(null, type);
this.owner = owner;
this.resultNumber = resultNumber;
}
};
_.Attribute = class {
};
_.TypedAttr = class extends _.Attribute {
constructor(value, type) {
super();
this.value = value;
this.type = type;
}
toString() {
return this.value;
}
};
_.StringAttr = class extends _.TypedAttr {
constructor(value, type) {
super(value, type || new _.PrimitiveType('string'));
}
toString() {
return this.value;
}
};
_.UnitAttr = class extends _.Attribute {
toString() {
return '';
}
};
_.IntegerAttr = class extends _.Attribute {
constructor(type, value) {
super();
this.value = value;
this.type = type;
}
toString() {
return this.value.toString();
}
};
_.BoolAttr = class extends _.Attribute {
constructor(value) {
super();
this.value = value;
this.type = new _.IntegerType('i1');
}
toString() {
return this.value ? 'true' : 'false';
}
};
_.FloatAttr = class extends _.Attribute {
constructor(type, value) {
super();
this.value = value;
this.type = type;
}
toString() {
return String(this.value);
}
};
_.MemRefLayoutAttr = class extends _.Attribute {
};
_.AffineMapAttr = class extends _.MemRefLayoutAttr {
constructor(map) {
super();
this._map = map;
}
get value() {
return this._map;
}
toString() {
return `affine_map<${this._map.toString()}>`;
}
};
_.IntegerSetAttr = class extends _.Attribute {
constructor(set) {
super();
this._set = set;
}
get value() {
return this._set;
}
toString() {
return `affine_set<${this._set.toString()}>`;
}
};
_.StridedLayoutAttr = class extends _.MemRefLayoutAttr {
constructor(offset, strides) {
super();
this.offset = offset;
this.strides = strides;
}
toString() {
const strides = `[${this.strides.join(', ')}]`;
if (this.offset !== null) {
return `strided<${strides}, offset: ${this.offset}>`;
}
return `strided<${strides}>`;
}
};
_.SymbolRefAttr = class extends _.Attribute {
constructor(rootReference, nestedReferences) {
super();
this.rootReference = rootReference;
this.nestedReferences = nestedReferences;
}
get value() {
if (this.nestedReferences && this.nestedReferences.length > 0) {
return `${this.rootReference}${this.nestedReferences.map((ref) => `::${ref}`).join('')}`;
}
return this.rootReference;
}
toString() {
return this.value;
}
};
_.DenseElementsAttr = class extends _.Attribute {
constructor(value, type) {
super();
this.value = value;
this.type = type;
}
};
_.SparseElementsAttr = class extends _.Attribute {
constructor(type, indices, values) {
super();
this.type = type;
this.indices = indices;
this.values = values;
}
};
_.DenseResourceElementsHandle = class {
constructor(key, blob = null) {
this.key = key;
this.blob = blob;
}
};
_.DenseResourceElementsAttr = class extends _.Attribute {
constructor(type, handle) {
super();
this.type = type;
this.rawHandle = handle;
}
get value() {
return this.rawHandle ? this.rawHandle.blob : null;
}
toString() {
const key = this.rawHandle ? this.rawHandle.key : 'unknown';
return `dense_resource<${key}>`;
}
};
_.OpaqueAttr = class extends _.Attribute {
constructor(dialectName, symbolData, type) {
super();
this.dialectName = dialectName;
this.symbolData = symbolData;
this.type = type;
}
get value() {
return this.toString();
}
toString() {
if (this.symbolData) {
return `${this.dialectName}.${this.symbolData}`;
}
return this.dialectName;
}
};
_.AffineLowPrecOp = {
LNoOp: '',
Add: 'Add',
Sub: 'Sub'
};
_.AffineHighPrecOp = {
HNoOp: '',
Mul: 'Mul',
FloorDiv: 'FloorDiv',
CeilDiv: 'CeilDiv',
Mod: 'Mod'
};
_.AffineExprKind = {
Add: 'Add',
Mul: 'Mul',
Mod: 'Mod',
FloorDiv: 'FloorDiv',
CeilDiv: 'CeilDiv',
Constant: 'Constant',
DimId: 'DimId',
SymbolId: 'SymbolId'
};
_.AffineExpr = class {
constructor(kind) {
this.kind = kind;
}
isSymbolicOrConstant() {
return this.kind === _.AffineExprKind.Constant || this.kind === _.AffineExprKind.SymbolId;
}
printAffineExprInternal(/* enclosingStrong */) {
return '';
}
toString() {
return this.printAffineExprInternal(false);
}
};
_.AffineDimExpr = class extends _.AffineExpr {
constructor(position) {
super(_.AffineExprKind.DimId);
this._position = position;
}
getPosition() {
return this._position;
}
toString() {
return `d${this._position}`;
}
};
_.AffineSymbolExpr = class extends _.AffineExpr {
constructor(position) {
super(_.AffineExprKind.SymbolId);
this._position = position;
}
getPosition() {
return this._position;
}
toString() {
return `s${this._position}`;
}
};
_.AffineConstantExpr = class extends _.AffineExpr {
constructor(value) {
super(_.AffineExprKind.Constant);
this.value = value;
}
toString() {
return String(this.value);
}
};
_.AffineBinaryOpExpr = class extends _.AffineExpr {
constructor(kind, lhs, rhs) {
super(kind);
this.lhs = lhs;
this.rhs = rhs;
}
isSymbolicOrConstant() {
return this.lhs.isSymbolicOrConstant() && this.rhs.isSymbolicOrConstant();
}
printAffineExprInternal(enclosingStrong) {
let op = '';
switch (this.kind) {
case _.AffineExprKind.Add: op = ' + '; break;
case _.AffineExprKind.Mul: op = ' * '; break;
case _.AffineExprKind.Mod: op = ' mod '; break;
case _.AffineExprKind.FloorDiv: op = ' floordiv '; break;
case _.AffineExprKind.CeilDiv: op = ' ceildiv '; break;
default: throw new Error(`Unexpected affine expression kind '${this.kind}'.`);
}
const isAdd = this.kind === _.AffineExprKind.Add;
const needsParens = enclosingStrong === true;
if (!isAdd) {
if (this.kind === _.AffineExprKind.Mul && this.rhs instanceof _.AffineConstantExpr && this.rhs.value === -1) {
const inner = `-${this.lhs.printAffineExprInternal(true)}`;
return needsParens ? `(${inner})` : inner;
}
const inner = `${this.lhs.printAffineExprInternal(true)}${op}${this.rhs.printAffineExprInternal(true)}`;
return needsParens ? `(${inner})` : inner;
}
if (this.rhs instanceof _.AffineBinaryOpExpr && this.rhs.kind === _.AffineExprKind.Mul &&
this.rhs.rhs instanceof _.AffineConstantExpr && this.rhs.rhs.value === -1) {
const lhsStr = this.lhs.printAffineExprInternal(false);
const rhsNeedsParens = this.rhs.lhs instanceof _.AffineBinaryOpExpr && this.rhs.lhs.kind === _.AffineExprKind.Add;
const rhsStr = this.rhs.lhs.printAffineExprInternal(rhsNeedsParens);
const inner = `${lhsStr} - ${rhsStr}`;
return needsParens ? `(${inner})` : inner;
}
if (this.rhs instanceof _.AffineConstantExpr && this.rhs.value < 0) {
const inner = `${this.lhs.printAffineExprInternal(false)} - ${-this.rhs.value}`;
return needsParens ? `(${inner})` : inner;
}
const inner = `${this.lhs.printAffineExprInternal(false)}${op}${this.rhs.printAffineExprInternal(false)}`;
return needsParens ? `(${inner})` : inner;
}
};
_.AffineMap = class {
constructor(numDims, numSymbols, results) {
this._numDims = numDims;
this._numSymbols = numSymbols;
this._results = results; // Array of AffineExpr or string for backward compat
}
getNumResults() {
return this._results.length;
}
getResults() {
return this._results;
}
toString() {
const dims = [];
for (let i = 0; i < this._numDims; i++) {
dims.push(`d${i}`);
}
const symbols = [];
for (let i = 0; i < this._numSymbols; i++) {
symbols.push(`s${i}`);
}
const dimsStr = `(${dims.join(', ')})`;
const symbolsStr = symbols.length > 0 ? `[${symbols.join(', ')}]` : '';
const resultsStr = this._results.map((r) => r.toString ? r.toString() : String(r)).join(', ');
return `${dimsStr}${symbolsStr} -> (${resultsStr})`;
}
static get(numDims, numSymbols, results) {
return new _.AffineMap(numDims, numSymbols, results);
}
// Backward compatible string-based parse
static parse(str) {
const arrowIdx = str.indexOf('->');
if (arrowIdx === -1) {
return new _.AffineMap(0, 0, [str]);
}
let resultPart = str.substring(arrowIdx + 2).trim();
const domainIdx = resultPart.indexOf(', domain:');
if (domainIdx !== -1) {
resultPart = resultPart.substring(0, domainIdx).trim();
}
// Parse result expressions from the result part (e.g., "(d0, d1 + s0)")
const results = [];
if (resultPart.startsWith('(')) {
let depth = 0;
let start = 1; // Skip opening paren
for (let i = 0; i < resultPart.length; i++) {
const ch = resultPart[i];
if (ch === '(' || ch === '[') {
depth++;
} else if (ch === ')' || ch === ']') {
depth--;
if (depth === 0) {
const expr = resultPart.substring(start, i).trim();
if (expr) {
results.push(expr);
}
break;
}
} else if (ch === ',' && depth === 1) {
const expr = resultPart.substring(start, i).trim();
if (expr) {
results.push(expr);
}
start = i + 1;
}
}
} else {
results.push(resultPart);
}
// For backward compat, store string representation of results
return new _.AffineMap(0, 0, results.length > 0 ? results : [str]);
}
};
// Integer set for affine_set constraints
_.IntegerSet = class {
constructor(numDims, numSymbols, constraints, eqFlags) {
this._numDims = numDims;
this._numSymbols = numSymbols;
this._constraints = constraints; // Array of AffineExpr
this._eqFlags = eqFlags; // Array of bool (true = equality, false = inequality)
}
getNumConstraints() {
return this._constraints.length;
}
getConstraints() {
return this._constraints;
}
getConstraint(idx) {
return this._constraints[idx];
}
getEqFlags() {
return this._eqFlags;
}
isEq(idx) {
return this._eqFlags[idx];
}
toString() {
const dims = [];
for (let i = 0; i < this._numDims; i++) {
dims.push(`d${i}`);
}
const symbols = [];
for (let i = 0; i < this._numSymbols; i++) {
symbols.push(`s${i}`);
}
const dimsStr = `(${dims.join(', ')})`;
const symbolsStr = symbols.length > 0 ? `[${symbols.join(', ')}]` : '';
const constraintsStr = this._constraints.map((c, i) => {
const exprStr = c.toString ? c.toString() : String(c);
return this._eqFlags[i] ? `${exprStr} == 0` : `${exprStr} >= 0`;
}).join(', ');
return `${dimsStr}${symbolsStr} : (${constraintsStr})`;
}
static get(numDims, numSymbols, constraints, eqFlags) {
return new _.IntegerSet(numDims, numSymbols, constraints, eqFlags);
}
};
_.IndexingMap = class {
constructor(affineMap) {
this._affineMap = affineMap;
}
GetAffineMap() {
return this._affineMap;
}
};
_.IndexingMapAttr = class extends _.Attribute {
constructor(indexingMap) {
super();
this._indexingMap = indexingMap;
}
get value() {
return this.toString();
}
getIndexingMap() {
return this._indexingMap;
}
static parse(parser /*, type */) {
parser.parseLess();
const indexingMap = _.IndexingMapAttr.parseChainOfStringsAsIndexingMap(parser);
parser.parseGreater();
return new _.IndexingMapAttr(indexingMap);
}
static parseChainOfStringsAsIndexingMap(parser) {
let indexingMapStr = '';
let str = parser.parseOptionalString();
while (str !== null) {
indexingMapStr += str;
str = parser.parseOptionalString();
}
return _.IndexingMapAttr.parseIndexingMap(indexingMapStr);
}
static parseIndexingMap(str) {
const domainIdx = str.indexOf(', domain:');
const mapStr = domainIdx >= 0 ? str.substring(0, domainIdx) : str.replace(/,$/, '');
const affineMap = _.AffineMap.parse(mapStr);
return new _.IndexingMap(affineMap);
}
toString() {
return `#xla.indexing_map<"${this._indexingMap.GetAffineMap().toString()}">`;
}
};
_.ArrayAttr = class extends _.Attribute {
constructor(elements) {
super();
this.elements = elements; // Array of Attribute objects
}
get value() {
return this.elements.map((e) => e && e.value !== undefined ? e.value : e);
}
toString() {
return `${this.elements.map((e) => e && e.toString ? e.toString() : String(e)).join(', ')}`;
}
};
_.DenseI64ArrayAttr = class extends _.Attribute {
constructor(values) {
super();
this._values = values; // Array of i64 integers
}
get value() {
return this._values;
}
toString() {
return `[${this._values.join(', ')}]`;
}
static parse(parser) {
if (!parser.parseOptionalLSquare()) {
return null;
}
const values = [];
while (!parser.parseOptionalRSquare()) {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
values.push(value);
parser.parseOptionalComma();
}
return new _.DenseI64ArrayAttr(values);
}
};
_.DictionaryAttr = class extends _.Attribute {
constructor(value) {
super();
this._value = value; // Map of name -> Attribute
}
get value() {
return this._value;
}
get(name) {
return this._value.get(name);
}
toString() {
const entries = Array.from(this._value.entries())
.map(([k, v]) => `${k} = ${v && v.toString ? v.toString() : String(v)}`);
return `{${entries.join(', ')}}`;
}
};
_.DenseArrayAttr = class extends _.Attribute {
constructor(type, size, data) {
super();
this.type = type;
this.size = size;
if (data instanceof Uint8Array) {
this.values = null;
this.data = data;
} else {
this.values = data;
}
}
get value() {
if (this.values === null) {
const blob = this.data;
const typeStr = this.type.toString();
const view = new DataView(blob.buffer, blob.byteOffset, blob.length);
const values = [];
if (typeStr.startsWith('i') || typeStr.startsWith('si') || typeStr.startsWith('ui') || typeStr === 'index') {
const match = typeStr.match(/[su]?i(\d+)/);
const bitWidth = match ? parseInt(match[1], 10) : 64;
const unsigned = typeStr.startsWith('ui');
const byteWidth = Math.ceil(bitWidth / 8);
const size = blob.length / byteWidth;
for (let i = 0; i < size; i++) {
if (bitWidth <= 8) {
values.push(unsigned ? view.getUint8(i * byteWidth) : view.getInt8(i * byteWidth));
} else if (bitWidth <= 16) {
values.push(unsigned ? view.getUint16(i * byteWidth, true) : view.getInt16(i * byteWidth, true));
} else if (bitWidth <= 32) {
values.push(unsigned ? view.getUint32(i * byteWidth, true) : view.getInt32(i * byteWidth, true));
} else {
values.push(unsigned ? view.getBigUint64(i * byteWidth, true) : view.getBigInt64(i * byteWidth, true));
}
}
} else if (typeStr === 'f32') {
const size = blob.length / 4;
for (let i = 0; i < size; i++) {
values.push(view.getFloat32(i * 4, true));
}
} else if (typeStr === 'f64') {
const size = blob.length / 8;
for (let i = 0; i < size; i++) {
values.push(view.getFloat64(i * 8, true));
}
} else if (typeStr === 'f16') {
const size = blob.length / 2;
for (let i = 0; i < size; i++) {
values.push(view.getFloat16(i * 2, true));
}
}
this.values = values;
this.data = null;
}
return this.values;
}
toString() {
return `array<${this.type.toString()}: ${this.value.join(', ')}>`;
}
};
_.DenseI32ArrayAttr = class extends _.DenseArrayAttr {
constructor(data) {
super(new _.IntegerType('i32'), Array.isArray(data) ? data.length : 0, data);
}
};
_.TypeAttrOf = class extends _.Attribute {
constructor(type) {
super();
this.type = type; // the type IS the value
}
toString() {
return this.type.toString();
}
};
_.ConvDimensionNumbersAttr = class extends _.Attribute {
constructor(input, kernel, output) {
super();
this.input = input;
this.kernel = kernel;
this.output = output;
}
toString() {
const formatDim = (dims) => `[${dims.join(', ')}]`;
return `${formatDim(this.input)}x${formatDim(this.kernel)}->${formatDim(this.output)}`;
}
};
_.StringRef = class {
constructor(value, position) {
this.value = value;
this.position = position;
}
data() {
return this.position;
}
str() {
return this.value;
}
drop_front(n = 1) {
return this.value.substring(n);
}
getAsInteger(radix) {
const str = this.value;
if (str.length === 0) {
return null;
}
const value = radix === 0 ? Number(str) : parseInt(str, radix);
if (!Number.isInteger(value)) {
return null;
}
return value;
}
toString() {
return this.value;
}
};
_.Type = class {
constructor(value) {
this._value = value;
}
get name() {
return this._value;
}
toString() {
return this._value;
}
};
_.OpaqueType = class extends _.Type {
constructor(dialect, symbolData) {
super();
this.dialect = dialect;
this.symbolData = symbolData;
}
toString() {
return `!${this.dialect}${this.symbolData}`;
}
};
_.util = {};
_.util.VariantType = class extends _.Type {
toString() {
return '?';
}
};
_.util.ListType = class extends _.Type {
constructor(elementType) {
super();
this.elementType = elementType;
}
static parse(parser) {
parser.parseLess();
const elementType = parser.parseOptionalQuestion()
? new _.util.VariantType()
: parser.parseType();
parser.parseGreater();
return new _.util.ListType(elementType);
}
toString() {
return `!util.list<${this.elementType}>`;
}
};
_.PrimitiveType = class extends _.Type {
};
_.IntegerType = class extends _.Type {
getWidth() {
const match = this._value.match(/^[su]?i(\d+)$/);
if (match) {
return parseInt(match[1], 10);
}
throw new mlir.Error(`Unexpected integer type '${this._value}'.`);
}
getIntOrFloatBitWidth() {
return this.getWidth();
}
};
_.IntegerType.kMaxWidth = (1 << 24) - 1;
_.FloatType = class extends _.Type {
getWidth() {
const name = this._value;
switch (name) {
case 'f16': return 16;
case 'f32': return 32;
case 'f64': return 64;
case 'f80': return 80;
case 'f128': return 128;
case 'bf16': return 16;
case 'tf32': return 32;
case 'f4E2M1FN': return 4;
case 'f6E2M3FN': return 6;
case 'f6E3M2FN': return 6;
case 'f8E3M4': return 8;
case 'f8E4M3': return 8;
case 'f8E4M3B11FNUZ': return 8;
case 'f8E4M3FN': return 8;
case 'f8E4M3FNUZ': return 8;
case 'f8E5M2': return 8;
case 'f8E5M2FNUZ': return 8;
case 'f8E8M0': return 8;
case 'f8E8M0FNU': return 8;
default: throw new mlir.Error(`Unexpected float type '${name}'.`);
}
}
getIntOrFloatBitWidth() {
return this.getWidth();
}
};
_.NoneType = class extends _.Type {
constructor() {
super('none');
}
};
_.TokenType = class extends _.Type {
constructor() {
super('token');
}
};
_.IndexType = class extends _.Type {
constructor() {
super('index');
}
};
_.FunctionType = class extends _.Type {
constructor(inputs, results) {
super(null);
this.inputs = inputs || [];
this.results = results || [];
}
toString() {
const inputs = this.inputs.map((t) => t.toString());
const results = this.results.map((t) => t.toString());
const result = results.length === 1 ? results[0] : `(${results.join(', ')})`;
return `(${inputs.join(', ')}) -> ${result}`;
}
};
_.ComplexType = class extends _.Type {
constructor(elementType) {
super(null);
this.elementType = elementType;
}
toString() {
const elementTypeStr = this.elementType?.toString ? this.elementType.toString() : this.elementType;
return `complex<${elementTypeStr}>`;
}
};
_.ShapedType = class extends _.Type {
getNumElements() {
if (this.shape.some((d) => d < 0n || d === _.ShapedType.kDynamic)) {
return 0n;
}
return this.shape.length === 0 ? 1n : this.shape.reduce((a, b) => a * b, 1n);
}
};
_.ShapedType.kDynamic = -9223372036854775808n;
_.RankedTensorType = class extends _.ShapedType {
constructor(shape, elementType, encoding) {
super(null);
this.shape = shape || [];
this.elementType = elementType;
this.encoding = encoding;
}
toString() {
const shapeStr = this.shape.map((d) => d < 0n ? '?' : d.toString()).join('x');
const elementTypeStr = this.elementType?.toString ? this.elementType.toString() : this.elementType;
const prefix = shapeStr ? `${shapeStr}x` : '';
if (this.encoding) {
return `tensor<${prefix}${elementTypeStr}, ${this.encoding}>`;
}
return `tensor<${prefix}${elementTypeStr}>`;
}
};
_.UnrankedTensorType = class extends _.Type {
constructor(elementType) {
super(null);
this.elementType = elementType;
}
toString() {
const elementTypeStr = this.elementType.toString();
return `tensor<*x${elementTypeStr}>`;
}
};
_.VectorType = class extends _.ShapedType {
constructor(shape, elementType, scalableDims) {
super(null);
this.shape = shape || [];
this.elementType = elementType;
this.scalableDims = scalableDims || [];
}
toString() {
const parts = this.shape.map((d, i) => {
const isScalable = this.scalableDims[i];
return isScalable ? `[${d}]` : String(d);
});
const shapeStr = parts.join('x');
const elementTypeStr = this.elementType?.toString ? this.elementType.toString() : this.elementType;
const prefix = shapeStr ? `${shapeStr}x` : '';
return `vector<${prefix}${elementTypeStr}>`;
}
};
_.MemRefType = class extends _.ShapedType {
constructor(shape, elementType, layout, memorySpace) {
super(null);
this.shape = shape || [];
this.elementType = elementType;
this.layout = layout || null;
this.memorySpace = memorySpace || null;
}
toString() {
const shapeStr = this.shape.map((d) => d < 0n ? '?' : d.toString()).join('x');
const elementTypeStr = this.elementType?.toString ? this.elementType.toString() : this.elementType;
const prefix = shapeStr ? `${shapeStr}x` : '';
let result = `memref<${prefix}${elementTypeStr}`;
if (this.layout) {
result += `, ${this.layout.toString()}`;
}
if (this.memorySpace) {
result += `, ${this.memorySpace.toString()}`;
}
result += '>';
return result;
}
};
_.UnrankedMemRefType = class extends _.Type {
constructor(elementType, memorySpace) {
super(null);
this.elementType = elementType;
this.memorySpace = memorySpace || null;
}
toString() {
const elementTypeStr = this.elementType?.toString ? this.elementType.toString() : this.elementType;
let result = `memref<*x${elementTypeStr}`;
if (this.memorySpace) {
result += `, ${this.memorySpace.toString()}`;
}
result += '>';
return result;
}
};
_.TupleType = class extends _.Type {
constructor(types) {
super(null);
this.types = types || [];
}
getTypes() {
return this.types;
}
getNumTypes() {
return this.types.length;
}
getType(index) {
return this.types[index];
}
toString() {
const typeStrs = this.types.map((t) => t?.toString ? t.toString() : t);
return `tuple<${typeStrs.join(', ')}>`;
}
};
_.SMLoc = class {
constructor(position) {
this.position = position || 0;
}
static getFromPointer(position) {
return new _.SMLoc(position);
}
copy() {
return new _.SMLoc(this.position);
}
};
_.Token = class {
constructor(decoder) {
this.loc = new _.SMLoc();
this.kind = null;
this.spelling = new _.StringRef('', decoder.position);
}
is(kind) {
return this.kind === kind;
}
isAny(...args) {
return args.some((kind) => this.kind === kind);
}
isNot(kind) {
return this.kind !== kind;
}
isKeyword() {
return this.kind.startsWith('kw_');
}
getSpelling() {
return this.spelling;
}
getUnsignedIntegerValue() {
const isHex = this.spelling.str().length > 1 && this.spelling.str()[1] === 'x';
const value = this.spelling.getAsInteger(isHex ? 0 : 10);
if (value === null || value < 0 || value > 0xFFFFFFFF) {
return null;
}
return value;
}
getUInt64IntegerValue() {
const spelling = this.spelling.str();
if (/^-?(\d+|0x[\da-f]+|0o[0-7]+|0b[01]+)$/i.test(spelling)) {
return BigInt(spelling);
}
return null;
}
getFloatingPointValue() {
const result = parseFloat(this.spelling.str());
if (isNaN(result)) {
return null;
}
return result;
}
getIntTypeBitwidth() {
const spelling = this.spelling.str();
const start = spelling[0] === 'i' ? 1 : 2;
const result = parseInt(spelling.slice(start), 10);
if (isNaN(result)) {
return null;
}
return result;
}
getIntTypeSignedness() {
const spelling = this.spelling.str();
if (spelling[0] === 'i') {
return null;
}
if (spelling[0] === 's') {
return true;
}
return false;
}
getHexStringValue() {
const spelling = this.spelling.str();
const bytes = spelling.slice(1, -1);
if (!bytes.startsWith('0x') || (bytes.length - 2) % 2 !== 0) {
return null;
}
const hex = bytes.slice(2);
const result = [];
for (let i = 0; i < hex.length; i += 2) {
const hi = parseInt(hex[i], 16);
const lo = parseInt(hex[i + 1], 16);
if (isNaN(hi) || isNaN(lo)) {
return null;
}
result.push((hi << 4) | lo);
}
return result;
}
getSymbolReference() {
const nameStr = this.spelling.str().slice(1);
if (nameStr[0] === '"') {
return this.getStringValue();
}
return nameStr;
}
getHashIdentifierNumber() {
const str = this.spelling.str().slice(1);
const result = parseInt(str, 10);
if (isNaN(result)) {
return null;
}
return result;
}
getStringValue() {
let bytes = this.spelling.str();
bytes = bytes.slice(1);
bytes = bytes.slice(0, -1);
if (this.kind === _.Token.at_identifier) {
bytes = bytes.slice(1);
}
if (!bytes.includes('\\')) {
return bytes;
}
const parts = [];
for (let i = 0; i < bytes.length;) {
const c = bytes[i++];
if (c !== '\\') {
parts.push(c);
continue;
}
const c1 = bytes[i++];
switch (c1) {
case '"':
case '\\':
parts.push(c1);
continue;
case 'n':
parts.push('\n');
continue;
case 't':
parts.push('\t');
continue;
default:
break;
}
const c2 = bytes[i++];
const hexValue = (parseInt(c1, 16) << 4) | parseInt(c2, 16);
parts.push(String.fromCharCode(hexValue));
}
return parts.join('');
}
};
_.Token.arrow = '->';
_.Token.at_identifier = '@';
_.Token.bare_identifier = 'id';
_.Token.caret_identifier = '^';
_.Token.colon = ':';
_.Token.comma = ',';
_.Token.ellipsis = 'ellipsis';
_.Token.eof = 'eof';
_.Token.equal = '=';
_.Token.exclamation_identifier = '!';
_.Token.file_metadata_begin = '{-#';
_.Token.file_metadata_end = '#-}';
_.Token.floatliteral = 'float';
_.Token.greater = '>';
_.Token.hash_identifier = '#';
_.Token.integer = 'int';
_.Token.inttype = 'inttype';
_.Token.kw_ceildiv = 'kw_ceildiv';
_.Token.kw_floordiv = 'kw_floordiv';
_.Token.kw_mod = 'kw_mod';
_.Token.l_brace = '{';
_.Token.l_paren = '(';
_.Token.l_square = '[';
_.Token.less = '<';
_.Token.minus = 'minus';
_.Token.percent_identifier = '%';
_.Token.plus = 'plus';
_.Token.question = '?';
_.Token.r_brace = '}';
_.Token.r_paren = ')';
_.Token.r_square = ']';
_.Token.slash = '/';
_.Token.star = '*';
_.Token.string = 'string';
_.Token.vertical_bar = '|';
_.Token.kw_affine_map = 'kw_affine_map';
_.Token.kw_affine_set = 'kw_affine_set';
_.Token.kw_array = 'kw_array';
_.Token.kw_dense = 'kw_dense';
_.Token.kw_dense_resource = 'kw_dense_resource';
_.Token.kw_distinct = 'kw_distinct';
_.Token.kw_false = 'kw_false';
_.Token.kw_loc = 'kw_loc';
_.Token.kw_offset = 'kw_offset';
_.Token.kw_sparse = 'kw_sparse';
_.Token.kw_strided = 'kw_strided';
_.Token.kw_symbol = 'kw_symbol';
_.Token.kw_true = 'kw_true';
_.Token.kw_unit = 'kw_unit';
_.MemoryBuffer = class {
constructor(decoder, identifier) {
this._decoder = decoder;
this._identifier = identifier || '';
}
get decoder() {
return this._decoder;
}
getBufferIdentifier() {
return this._identifier;
}
};
_.SrcBuffer = class {
constructor(buffer, includeLoc) {
this._buffer = buffer;
this._includeLoc = includeLoc || new _.SMLoc();
this._offsetCache = null;
}
get buffer() {
return this._buffer;
}
get includeLoc() {
return this._includeLoc;
}
getLineNumber(pointer) {
const offsets = this._getOrCreateOffsetCache();
let low = 0;
let high = offsets.length - 1;
while (low < high) {
const mid = (low + high + 1) >> 1;
if (offsets[mid] <= pointer) {
low = mid;
} else {
high = mid - 1;
}
}
return low + 1;
}
getPointerForLineNumber(line) {
const offsets = this._getOrCreateOffsetCache();
return offsets[line - 1];
}
_getOrCreateOffsetCache() {
if (!this._offsetCache) {
const buffer = this._buffer.decoder.buffer;
const offsets = [0];
if (buffer) {
for (let i = 0; i < buffer.length; i++) {
if (buffer[i] === 0x0A) {
offsets.push(i + 1);
}
}
}
this._offsetCache = offsets;
}
return this._offsetCache;
}
};
_.SourceMgr = class {
constructor() {
this._buffers = [];
}
isValidBufferID(id) {
return id > 0 && id <= this._buffers.length;
}
addNewSourceBuffer(buffer, includeLoc) {
this._buffers.push(new _.SrcBuffer(buffer, includeLoc));
return this._buffers.length;
}
getNumBuffers() {
return this._buffers.length;
}
getMainFileID() {
return 1;
}
getBufferInfo(id) {
return this._buffers[id - 1];
}
getMemoryBuffer(id) {
return this._buffers[id - 1].buffer;
}
getParentIncludeLoc(id) {
return this._buffers[id - 1].includeLoc;
}
};
_.Lexer = class {
constructor(sourceMgr, context) {
this._sourceMgr = sourceMgr;
this._context = context;
const bufferID = sourceMgr.getMainFileID();
this._decoder = sourceMgr.getMemoryBuffer(bufferID).decoder;
this._currentPosition = this._decoder.position;
this._current = this._decoder.decode();
this._nextPosition = this._decoder.position;
this._next = this._decoder.decode();
this._tokens = [new _.Token(this._decoder), new _.Token(this._decoder), new _.Token(this._decoder), new _.Token(this._decoder)];
this._index = 0;
this._errorLoc = new _.SMLoc();
}
getSourceMgr() {
return this._sourceMgr;
}
get decoder() {
return this._decoder;
}
getEncodedSourceLocation(loc) {
const mainFileID = this._sourceMgr.getMainFileID();
const bufferInfo = this._sourceMgr.getBufferInfo(mainFileID);
const lineNo = bufferInfo.getLineNumber(loc.position);
const column = (loc.position - bufferInfo.getPointerForLineNumber(lineNo)) + 1;
const identifier = this._sourceMgr.getMemoryBuffer(mainFileID).getBufferIdentifier();
return new _.Location(new _.FileLineColLoc(this._context, identifier, lineNo, column));
}
emitError(loc, message) {
const location = this.getEncodedSourceLocation(_.SMLoc.getFromPointer(loc));
throw new mlir.Error(`${message} at ${location.impl.printLocationInternal()}.`);
}
lexToken() {
this._position = this._currentPosition;
while (this._current) {
switch (this._current) {
case ' ':
case '\t':
case '\n':
case '\r':
this._skipWhitespace();
this._position = this._currentPosition;
continue;
case '/':
if (this._peek() !== '/') {
this._read();
return this.formToken('/', '/');
}
this.lexComment();
this._position = this._currentPosition;
continue;
case '.':
this._read();
if (this._current === '.' && this._next === '.') {
this._read();
this._read();
return this.formToken('ellipsis', '...');
}
this.emitError(this._position, 'Expected three consecutive dots for an ellipsis');
continue;
case '-':
if (this._peek() === '>') {
this._read();
this._read();
return this.formToken('->', '->');
}
this._read();
return this.formToken('minus', '-');
case '+':
this._read();
return this.formToken('plus', '+');
case '"':
return this.lexString();
case '@':
return this.lexPrefixedIdentifier('@');
case '%':
return this.lexPrefixedIdentifier('%');
case '#':
if (this._peek() === '-') {
const position = this._decoder.position;
const next = this._decoder.decode();
this._decoder.position = position;
if (next === '}') {
this._read();
this._read();
this._read();
return this.formToken('#-}', '#-}');
}
}
return this.lexPrefixedIdentifier('#');
case '!':
return this.lexPrefixedIdentifier('!');
case '^':
return this.lexPrefixedIdentifier('^');
case '=':
this._read();
return this.formToken('=', '=');
case ':':
this._read();
return this.formToken(':', ':');
case ',':
case '(':
case ')':
case '}':
case '[':
case ']':
case '<':
case '>':
case '?':
case '*':
case '|': {
const value = this._read();
return this.formToken(value, value);
}
case '{':
if (this._peek() === '-') {
const position = this._decoder.position;
const next = this._decoder.decode();
this._decoder.position = position;
if (next === '#') {
this._read();
this._read();
this._read();
return this.formToken('{-#', '{-#');
}
}
this._read();
return this.formToken('{', '{');
default:
if (/[a-zA-Z_]/.test(this._current)) {
return this.lexBareIdentifierOrKeyword();
}
if (/[0-9]/.test(this._current)) {
return this.lexNumber();
}
this.emitError(this._currentPosition, 'Unexpected character');
}
}
return this.formToken('eof', null);
}
resetPointer(newPointer) {
if (newPointer < 0) {
throw new mlir.Error('Invalid negative offset.');
}
this._decoder.position = newPointer;
this._nextPosition = this._decoder.position;
this._next = this._decoder.decode();
this._read();
}
_read() {
const current = this._current;
this._current = this._next;
this._currentPosition = this._nextPosition;
this._nextPosition = this._decoder.position;
this._next = this._decoder.decode();
return current;
}
_peek() {
return this._next;
}
_eat(value) {
if (this._current === value) {
this._read();
return true;
}
return false;
}
_skipWhitespace() {
while (this._current !== undefined && (this._current === ' ' || this._current === '\t' || this._current === '\n' || this._current === '\r')) {
this._read();
}
}
lexComment() {
this._read('/');
if (this._current !== '/') {
this.emitError(this._position, 'Unexpected character');
}
while (this._current && this._current !== '\n' && this._current !== '\r') {
this._read();
}
}
lexNumber() {
let v = '';
let type = 'int';
while (this._current && /[0-9]/.test(this._current)) {
v += this._read();
}
if (v === '0' && this._current === 'x' && /[0-9a-fA-F]/.test(this._peek())) {
v += this._read();
while (this._current && /[0-9a-fA-F]/.test(this._current)) {
v += this._read();
}
return this.formToken(type, v);
}
if (this._current === '.') {
v += this._read();
type = 'float';
while (this._current && /[0-9]/.test(this._current)) {
v += this._read();
}
if (this._current === 'e' || this._current === 'E') {
const next1 = this._peek();
let next2 = '';
if (next1 === '+' || next1 === '-') {
const position = this._decoder.position;
next2 = this._decoder.decode();
this._decoder.position = position;
}
if (/[0-9]/.test(next1) || ((next1 === '+' || next1 === '-') && /[0-9]/.test(next2))) {
v += this._read();
if (this._current === '+' || this._current === '-') {
v += this._read();
}
while (this._current && /[0-9]/.test(this._current)) {
v += this._read();
}
}
}
return this.formToken(type, v);
}
return this.formToken(type, v);
}
lexString() {
const parts = ['"'];
this._read();
let chunk = '';
while (this._current && this._current !== '"') {
if (this._current === '\n' || this._current === '\v' ||
this._current === '\f' || this._current === '\r') {
this.emitError(this._position, 'Expected \'"\' in string literal');
}
if (this._current === '\\') {
if (chunk) {
parts.push(chunk);
chunk = '';
}
parts.push(this._current);
this._read();
if (this._current === '"' || this._current === '\\' || this._current === 'n' || this._current === 't') {
parts.push(this._current);
this._read();
} else if (this._current && /[0-9a-fA-F]/.test(this._current) && this._next && /[0-9a-fA-F]/.test(this._next)) {
parts.push(this._current);
this._read();
parts.push(this._current);
this._read();
} else {
this.emitError(this._position, 'Unknown escape in string literal');
}
} else {
chunk += this._current;
if (chunk.length >= 4096) {
parts.push(chunk);
chunk = '';
}
this._read();
}
}
if (chunk) {
parts.push(chunk);
}
if (this._eat('"')) {
parts.push('"');
return this.formToken(_.Token.string, parts.join(''));
}
this.emitError(this._position, "expected '\"' in string literal");
return null;
}
lexBareIdentifierOrKeyword() {
let result = '';
while (this._current && (/[a-zA-Z_$.]/.test(this._current) || /[0-9]/.test(this._current))) {
result += this._read();
}
const isAllDigit = (str) => /^[0-9]+$/.test(str);
if ((result.length > 1 && result[0] === 'i' && isAllDigit(result.slice(1))) ||
(result.length > 2 && result[1] === 'i' && (result[0] === 's' || result[0] === 'u') && isAllDigit(result.slice(2)))) {
return this.formToken('inttype', result);
}
switch (result) {
case 'affine_map':
case 'affine_set':
case 'array':
case 'ceildiv':
case 'dense':
case 'dense_resource':
case 'distinct':
case 'floordiv':
case 'loc':
case 'mod':
case 'offset':
case 'sparse':
case 'strided':
case 'symbol':
case 'unit':
return this.formToken(`kw_${result}`, result);
case 'true':
return this.formToken(_.Token.kw_true, result);
case 'false':
return this.formToken(_.Token.kw_false, result);
default:
return this.formToken(_.Token.bare_identifier, result);
}
}
lexPrefixedIdentifier(prefix) {
let result = prefix;
this._read();
if (prefix === '@' && this._current === '"') {
result += this.lexString().getSpelling().str();
return this.formToken(prefix, result);
}
if (prefix === '@') {
// symbol-ref-id ::= `@` (bare-id | string-literal)
// bare-id ::= (letter|[_]) (letter|digit|[_$.])*
if (this._current && /[a-zA-Z_]/.test(this._current)) {
while (this._current && /[a-zA-Z0-9_$.]/.test(this._current)) {
result += this._read();
}
} else if (result.length === 1) {
this.emitError(this._position, '@ identifier expected to start with letter or \'_\'');
}
} else if (this._current && /[0-9]/.test(this._current)) {
// suffix-id ::= digit+ | ...
while (this._current && /[0-9]/.test(this._current)) {
result += this._read();
}
} else if (this._current && /[a-zA-Z_$.-]/.test(this._current)) {
// suffix-id ::= ... | (letter|id-punct) (letter|id-punct|digit)*
while (this._current && /[a-zA-Z_$0-9.-]/.test(this._current)) {
if (this._current === '-' && this._peek() === '>') {
break;
}
result += this._read();
}
} else if (result.length === 1) {
const errorKinds = { '#': 'Invalid attribute name', '%': 'Invalid SSA name', '^': 'Invalid block name', '!': 'Invalid type identifier' };
this.emitError(this._position, errorKinds[prefix] || 'Invalid identifier');
}
if (prefix === '@' && this._current === ':' && this._peek() === ':') {
result += this._read();
result += this._read();
result += this.lexPrefixedIdentifier('@').getSpelling().str();
}
const kind = prefix === '$' ? '%' : prefix;
return this.formToken(kind, result);
}
formToken(kind, value) {
const token = this._tokens[this._index];
this._index = (this._index + 1) % this._tokens.length;
token.loc.position = this._position;
token.kind = kind;
token.spelling.value = value;
token.spelling.position = this._position;
return token;
}
};
_.AsmResourceParser = class {
constructor(name) {
this.name = name;
}
};
_.ParsedResourceEntry = {};
_.ParsedResourceEntry.Text = class {
constructor(key, keyLoc, value, parser) {
this.key = key;
this.keyLoc = keyLoc;
this.value = value;
this.parser = parser;
}
parseAsBool() {
if (this.value.kind === 'boolean') {
return this.value.value;
}
throw new mlir.Error(`Expected boolean value for resource entry '${this.key}'`);
}
parseAsString() {
if (this.value.kind === _.Token.string) {
return this.value.getStringValue();
}
throw new mlir.Error(`Expected string value for resource entry '${this.key}'`);
}
parseAsBlob() {
if (this.value.kind === _.Token.string) {
const hexStr = this.value.getStringValue();
if (hexStr.startsWith('0x')) {
const hex = hexStr.slice(2);
const bytes = new Uint8Array(hex.length / 2);
for (let i = 0; i < bytes.length; i++) {
bytes[i] = parseInt(hex.slice(i * 2, i * 2 + 2), 16);
}
if (bytes.length >= 4) {
return bytes.slice(4);
}
return bytes;
}
}
throw new mlir.Error(`Expected hex string blob value for resource entry '${this.key}'`);
}
};
_.ParserConfig = class {
constructor(context) {
this.context = context;
this.resourceParsers = new Map();
}
getResourceParser(name) {
return this.resourceParsers.get(name);
}
attachResourceParser(parser) {
this.resourceParsers.set(parser.name, parser);
}
};
_.ParserState = class {
constructor(sourceMgr, config) {
this.config = config;
this.defaultDialectStack = ['builtin'];
this.attributeAliasDefinitions = new Map();
this.typeAliasDefinitions = new Map();
this.dialectResources = new Map();
this.deferredLocsReferences = [];
this.lex = new _.Lexer(sourceMgr, config.context);
this.curToken = this.lex.lexToken();
}
};
_.IsolatedSSANameScope = class {
constructor() {
this.values = new Map();
this.definitionsPerScope = [];
}
recordDefinition(def) {
this.definitionsPerScope[this.definitionsPerScope.length - 1].add(def);
}
pushSSANameScope() {
this.definitionsPerScope.push(new Set());
}
popSSANameScope() {
for (const def of this.definitionsPerScope.pop()) {
this.values.delete(def);
}
}
};
_.Parser = class {
constructor(state) {
this.state = state;
}
get context() {
return this.state.config.context;
}
parseCommaSeparatedListUntil(rightToken, parseElement, allowEmptyList = true) {
if (this.getToken().is(rightToken)) {
if (!allowEmptyList) {
this.emitError('Expected list element');
}
this.consumeToken(rightToken);
return;
}
this.parseCommaSeparatedList('none', parseElement);
this.parseToken(rightToken, `Expected '${rightToken}'`);
}
parseResourceFileMetadata(parseBody) {
this.parseToken(_.Token.l_brace, "Expected '{'");
this.parseCommaSeparatedListUntil(_.Token.r_brace, () => {
const nameLoc = this.getToken().loc.copy();
const name = this.parseOptionalKeyword();
if (!name) {
this.emitError("Expected identifier key for 'resource' entry");
}
this.parseToken(_.Token.colon);
this.parseToken(_.Token.l_brace);
parseBody(name, nameLoc);
});
}
parseDialectResourceFileMetadata() {
this.parseResourceFileMetadata((name, nameLoc) => {
const dialect = this.context.getOrLoadDialect(name);
if (!dialect) {
this.emitError(nameLoc, `Dialect '${name}' is unknown`);
}
this.parseCommaSeparatedListUntil(_.Token.r_brace, () => {
const keyLoc = this.getToken().loc.copy();
const handle = this.parseResourceHandle(dialect);
const key = dialect.getResourceKey(handle);
this.parseToken(_.Token.colon, "Expected ':'");
const valueTok = this.getToken();
this.consumeToken();
const entry = new _.ParsedResourceEntry.Text(key, keyLoc, valueTok, this);
dialect.parseResource(entry);
});
});
}
parseExternalResourceFileMetadata() {
this.parseResourceFileMetadata((name, nameLoc) => {
const handler = this.state.config.getResourceParser(name);
if (!handler) {
this.emitWarning(this.getEncodedSourceLocation(nameLoc), `Ignoring unknown external resources for '${name}'`);
}
this.parseCommaSeparatedListUntil(_.Token.r_brace, () => {
const keyLoc = this.getToken().loc;
const key = this.parseOptionalKeywordOrString();
if (!key) {
this.emitError("Expected identifier key for 'external_resources' entry");
}
this.parseToken(_.Token.colon, "Expected ':'");
const valueTok = this.getToken();
this.consumeToken();
if (handler && handler.parseResource) {
const entry = new _.ParsedResourceEntry.Text(key, keyLoc, valueTok, this);
handler.parseResource(entry);
}
});
});
}
parseOptionalKeywordOrString() {
const keyword = this.parseOptionalKeyword();
if (keyword) {
return keyword;
}
return this.parseOptionalString();
}
finalize(/* block */) {
}
isCurrentTokenAKeyword() {
return this.getToken().isAny(_.Token.bare_identifier, _.Token.inttype) || this.getToken().isKeyword();
}
parseOptionalKeyword() {
if (!this.isCurrentTokenAKeyword()) {
return null;
}
const keyword = this.getTokenSpelling().str();
this.consumeToken();
return keyword;
}
parseTypeListNoParens() {
return this.parseCommaSeparatedList('none', () => this.parseType());
}
parseTypeListParens() {
this.parseToken(_.Token.l_paren, "Expected '('");
if (this.consumeIf(_.Token.r_paren)) {
return [];
}
const types = this.parseTypeListNoParens();
this.parseToken(_.Token.r_paren, "Expected ')'");
return types;
}
parseSuccessors(successors) {
const parsed = this.parseCommaSeparatedList('square', () => {
const label = this.getTokenSpelling().str();
this.consumeToken(_.Token.caret_identifier);
return { label };
});
if (parsed.length === 0) {
this.emitError('Expected at least one successor');
}
for (const s of parsed) {
successors.push(s);
}
}
parseAttributeDict(attributes) {
const seenKeys = new Set();
this.parseCommaSeparatedList('brace', () => {
// The name of an attribute can either be a bare identifier, or a string.
let name = null;
if (this.getToken().is(_.Token.string)) {
name = this.getToken().getStringValue();
} else if (this.getToken().isAny(_.Token.bare_identifier, _.Token.inttype) || this.getToken().isKeyword()) {
name = this.getTokenSpelling().str();
} else {
this.emitError('Expected attribute name');
}
if (!name) {
this.emitError('Expected attribute name');
}
this.consumeToken();
if (seenKeys.has(name)) {
this.emitError(`Duplicate key '${name}' in dictionary attribute`);
}
seenKeys.add(name);
if (!this.consumeIf(_.Token.equal)) {
attributes.set(name, new _.UnitAttr());
return;
}
const attr = this.parseAttribute();
attributes.set(name, attr);
});
}
parseLocationInstance() {
if (this.getToken().is(_.Token.hash_identifier)) {
const locAttr = this.parseExtendedAttr();
if (locAttr instanceof _.LocationAttr === false && locAttr instanceof _.OpaqueAttr === false) {
this.emitError(`Expected location attribute, but got '${locAttr}'`);
}
return locAttr;
}
if (this.getToken().is(_.Token.string)) {
return this.parseNameOrFileLineColRange();
}
if (!this.getToken().is(_.Token.bare_identifier)) {
this.emitError(`Expected location instance, got '${this.getTokenSpelling().str()}'`);
}
if (this.getToken().spelling.str() === 'callsite') {
return this.parseCallSiteLocation();
}
if (this.getToken().spelling.str() === 'fused') {
return this.parseFusedLocation();
}
if (this.getToken().spelling.str() === 'unknown') {
this.consumeToken(_.Token.bare_identifier);
return new _.UnknownLoc();
}
this.emitError(`Expected location instance, got '${this.getTokenSpelling().str()}'`);
return null;
}
parseCallSiteLocation() {
this.consumeToken(_.Token.bare_identifier);
this.parseToken(_.Token.l_paren, "Expected '(' in callsite location");
const callee = this.parseLocationInstance();
if (this.getToken().isNot(_.Token.bare_identifier) || this.getToken().getSpelling().str() !== 'at') {
this.emitError("Expected 'at' in callsite location");
}
this.consumeToken(_.Token.bare_identifier);
const caller = this.parseLocationInstance();
this.parseToken(_.Token.r_paren, "Expected ')' in callsite location");
return new _.CallSiteLoc(callee, caller);
}
parseFusedLocation() {
this.consumeToken(_.Token.bare_identifier);
let metadata = null;
if (this.consumeIf(_.Token.less)) {
metadata = this.parseAttribute();
if (!metadata) {
this.emitError('Expected attribute in fused location metadata');
}
this.parseToken(_.Token.greater, "Expected '>' after fused location metadata");
}
const locations = this.parseCommaSeparatedList('square', () => this.parseLocationInstance());
return new _.FusedLoc(locations, metadata);
}
parseNameOrFileLineColRange() {
const str = this.getToken().getStringValue();
this.consumeToken(_.Token.string);
if (this.consumeIf(_.Token.colon)) {
if (this.getToken().isNot(_.Token.integer)) {
this.emitError('Expected integer line number in FileLineColRange');
}
const startLine = this.getToken().getUnsignedIntegerValue();
if (startLine === null) {
this.emitError('Expected integer line number in FileLineColRange');
}
this.consumeToken(_.Token.integer);
if (this.getToken().isNot(_.Token.colon)) {
return new _.FileLineColRange(str, startLine);
}
this.consumeToken(_.Token.colon);
if (this.getToken().isNot(_.Token.integer)) {
this.emitError('Expected integer column number in FileLineColRange');
}
const startCol = this.getToken().getUnsignedIntegerValue();
if (startCol === null) {
this.emitError('Expected integer column number in FileLineColRange');
}
this.consumeToken(_.Token.integer);
if (!this.isCurrentTokenAKeyword() || this.getTokenSpelling().str() !== 'to') {
return new _.FileLineColLoc(str, startLine, startCol);
}
this.consumeToken();
let endLine = null;
if (this.getToken().is(_.Token.integer)) {
endLine = this.getToken().getUnsignedIntegerValue();
if (endLine === null) {
this.emitError('Expected integer line number in FileLineColRange');
}
this.consumeToken(_.Token.integer);
}
if (this.getToken().isNot(_.Token.colon)) {
this.emitError("Expected either integer or ':' post 'to' in FileLineColRange");
}
this.consumeToken(_.Token.colon);
if (this.getToken().isNot(_.Token.integer)) {
this.emitError('Expected integer column number in FileLineColRange');
}
const endCol = this.getToken().getUnsignedIntegerValue();
if (endCol === null) {
this.emitError('Expected integer column number in FileLineColRange');
}
this.consumeToken(_.Token.integer);
if (endLine !== null) {
return new _.FileLineColRange(str, startLine, startCol, endLine, endCol);
}
return new _.FileLineColRange(str, startLine, startCol, undefined, endCol);
}
if (this.consumeIf(_.Token.l_paren)) {
const childLoc = this.parseLocationInstance();
this.parseToken(_.Token.r_paren, "Expected ')' after child location of NameLoc");
return new _.NameLoc(str, childLoc);
}
return new _.NameLoc(str);
}
parseLocationAlias() {
const tok = this.getToken();
this.consumeToken(_.Token.hash_identifier);
const identifier = tok.getSpelling().str().substring(1); // drop_front - remove '#' prefix
const attr = this.state.attributeAliasDefinitions.get(tok.getSpelling().str());
if (attr) {
if (attr instanceof _.LocationAttr) {
return attr;
}
return attr;
}
const index = this.state.deferredLocsReferences.length;
this.state.deferredLocsReferences.push({ loc: tok.loc, identifier });
return new _.OpaqueLoc(index, identifier, new _.UnknownLoc());
}
parseOptionalLocationSpecifier() {
if (!this.consumeIf(_.Token.kw_loc)) {
return null;
}
this.parseToken(_.Token.l_paren, "expected '(' in location");
const tok = this.getToken();
let directLoc = null;
if (tok.is(_.Token.hash_identifier) && !tok.getSpelling().str().includes('.')) {
directLoc = this.parseLocationAlias();
} else {
directLoc = this.parseLocationInstance();
}
this.parseToken(_.Token.r_paren, "expected ')' in location");
return directLoc;
}
parseXInDimensionList() {
// Matches ref impl: if token is 'xf32', reset lexer to after 'x'
if (this.getToken().isNot(_.Token.bare_identifier)) {
return false;
}
const token = this.getToken().getSpelling().str();
if (!token.startsWith('x')) {
return false;
}
// If we had a prefix of 'x' (e.g., 'xf32'), reset lexer to after the 'x'
if (token.length > 1) {
this.state.lex.resetPointer(this.getToken().loc.position + 1);
}
this.state.curToken = this.state.lex.lexToken();
return true;
}
parseIntegerInDimensionList() {
const spelling = this.getTokenSpelling().str();
if (spelling[0] === '0' && spelling.length > 1 && spelling[1] === 'x') {
this.state.lex.resetPointer(this.getToken().loc.position + 1);
this.state.curToken = this.state.lex.lexToken();
return 0n;
}
const dimension = this.getToken().getUInt64IntegerValue();
if (dimension === null) {
this.emitError('Invalid dimension');
}
this.consumeToken(_.Token.integer);
return dimension;
}
parseVectorDimensionList() {
const dimensions = [];
const scalableDims = [];
while (this.getToken().is(_.Token.integer) || this.getToken().is(_.Token.l_square)) {
const scalable = this.consumeIf(_.Token.l_square);
dimensions.push(this.parseIntegerInDimensionList());
if (scalable) {
if (!this.consumeIf(_.Token.r_square)) {
this.emitError("Missing ']' closing scalable dimension");
}
}
scalableDims.push(scalable);
if (!this.parseXInDimensionList()) {
break;
}
}
return { dimensions, scalableDims };
}
parseDimensionListRanked(allowDynamic = true, withTrailingX = true) {
const dimensions = [];
const parseDim = () => {
if (allowDynamic && this.consumeIf(_.Token.question)) {
dimensions.push(_.ShapedType.kDynamic);
return true;
}
if (this.getToken().is(_.Token.integer)) {
dimensions.push(this.parseIntegerInDimensionList());
return true;
}
return false;
};
const hasDimToken = () => this.getToken().isAny(_.Token.integer, _.Token.question);
if (withTrailingX) {
while (hasDimToken()) {
if (!parseDim() || !this.parseXInDimensionList()) {
break;
}
}
} else if (parseDim()) {
while (this.getToken().is(_.Token.bare_identifier) && this.getTokenSpelling().str().startsWith('x')) {
if (!this.parseXInDimensionList() || !parseDim()) {
break;
}
}
}
return { dimensions };
}
parseTensorType() {
this.parseToken(_.Token.less, "Expected '<' in tensor type");
let isUnranked = false;
let dimensions = [];
if (this.consumeIf(_.Token.star)) {
isUnranked = true;
this.parseXInDimensionList();
} else {
const dimInfo = this.parseDimensionListRanked();
dimensions = dimInfo.dimensions;
}
const elementTypeLoc = this.getToken().loc.copy();
const elementType = this.parseType();
let encoding = null;
if (this.consumeIf(_.Token.comma)) {
encoding = this.parseAttribute();
}
this.parseToken(_.Token.greater, "Expected '>' in tensor type");
if (elementType instanceof _.NoneType || elementType instanceof _.FunctionType ||
elementType instanceof _.TupleType || elementType instanceof _.RankedTensorType ||
elementType instanceof _.UnrankedTensorType || elementType instanceof _.MemRefType ||
elementType instanceof _.UnrankedMemRefType) {
this.emitError(elementTypeLoc, 'Invalid tensor element type');
}
if (isUnranked) {
if (encoding) {
this.emitError('Cannot apply encoding to unranked tensor');
}
return new _.UnrankedTensorType(elementType);
}
return new _.RankedTensorType(dimensions, elementType, encoding);
}
parseMemRefType() {
this.parseToken(_.Token.less, "Expected '<' in memref type");
let isUnranked = false;
let dimensions = [];
if (this.consumeIf(_.Token.star)) {
isUnranked = true;
this.parseXInDimensionList();
} else {
const dimInfo = this.parseDimensionListRanked();
dimensions = dimInfo.dimensions;
}
const typeLoc = this.getToken().loc.copy();
const elementType = this.parseType();
if (elementType instanceof _.NoneType || elementType instanceof _.FunctionType ||
elementType instanceof _.TupleType || elementType instanceof _.RankedTensorType ||
elementType instanceof _.UnrankedTensorType) {
this.emitError(typeLoc, 'Invalid memref element type');
}
let layout = null;
let memorySpace = null;
while (this.consumeIf(_.Token.comma)) {
const attr = this.parseAttribute();
if (attr instanceof _.MemRefLayoutAttr) {
layout = attr;
if (isUnranked) {
this.emitError('Cannot have affine map for unranked memref type');
}
if (memorySpace) {
this.emitError('Expected memory space to be last in memref type');
}
} else {
if (memorySpace) {
this.emitError('Multiple memory spaces specified in memref type');
}
memorySpace = attr;
}
}
this.parseToken(_.Token.greater, "Expected '>' in memref type");
if (isUnranked) {
return new _.UnrankedMemRefType(elementType, memorySpace);
}
return new _.MemRefType(dimensions, elementType, layout, memorySpace);
}
parseVectorType() {
this.parseToken(_.Token.less, "Expected '<' in vector type");
const dimInfo = this.parseVectorDimensionList();
const elementType = this.parseType();
this.parseToken(_.Token.greater, "Expected '>' in vector type");
return new _.VectorType(dimInfo.dimensions, elementType, dimInfo.scalableDims);
}
parseComplexType() {
this.parseToken(_.Token.less, "Expected '<' in complex type");
const elementTypeLoc = this.getToken().loc.copy();
const elementType = this.parseType();
if (!(elementType instanceof _.FloatType) && !(elementType instanceof _.IntegerType)) {
this.emitError(elementTypeLoc, 'Invalid element type for complex');
}
this.parseToken(_.Token.greater, "Expected '>' in complex type");
return new _.ComplexType(elementType);
}
parseTupleType() {
this.parseToken(_.Token.less, "Expected '<' in tuple type");
if (this.consumeIf(_.Token.greater)) {
return new _.TupleType([]);
}
const types = this.parseCommaSeparatedList('none', () => this.parseType());
this.parseToken(_.Token.greater, "Expected '>' in tuple type");
return new _.TupleType(types);
}
parseType() {
if (this.getToken().is(_.Token.l_paren)) {
return this.parseFunctionType();
}
return this.parseNonFunctionType();
}
parseOptionalType() {
switch (this.getToken().kind) {
case _.Token.l_paren:
case _.Token.exclamation_identifier:
case _.Token.inttype:
return this.parseType();
case _.Token.bare_identifier: {
switch (this.getToken().getSpelling().str()) {
case 'memref':
case 'tensor':
case 'complex':
case 'tuple':
case 'vector':
case 'f4E2M1FN':
case 'f6E2M3FN':
case 'f6E3M2FN':
case 'f8E5M2':
case 'f8E4M3':
case 'f8E4M3FN':
case 'f8E5M2FNUZ':
case 'f8E4M3FNUZ':
case 'f8E4M3B11FNUZ':
case 'f8E3M4':
case 'f8E8M0FNU':
case 'bf16':
case 'f16':
case 'tf32':
case 'f32':
case 'f64':
case 'f80':
case 'f128':
case 'index':
case 'none':
case 'token':
return this.parseType();
default:
break;
}
break;
}
default: {
break;
}
}
return null;
}
parseNonFunctionType() {
switch (this.getToken().kind) {
case _.Token.inttype: {
const width = this.getToken().getIntTypeBitwidth();
if (width === null) {
this.emitError('Invalid integer width');
}
if (width > _.IntegerType.kMaxWidth) {
this.emitError(`Integer bitwidth is limited to ${_.IntegerType.kMaxWidth} bits`);
}
const signedness = this.getToken().getIntTypeSignedness();
this.consumeToken(_.Token.inttype);
let prefix = 'i';
if (signedness === true) {
prefix = 'si';
} else if (signedness === false) {
prefix = 'ui';
}
return new _.IntegerType(`${prefix}${width}`);
}
case _.Token.exclamation_identifier: {
return this.parseExtendedType();
}
case _.Token.bare_identifier: {
const value = this.getTokenSpelling().str();
this.consumeToken(_.Token.bare_identifier);
switch (value) {
case 'tensor':
return this.parseTensorType();
case 'vector':
return this.parseVectorType();
case 'memref':
return this.parseMemRefType();
case 'complex':
return this.parseComplexType();
case 'tuple':
return this.parseTupleType();
case 'none':
return new _.NoneType();
case 'token':
return new _.TokenType();
case 'index':
return new _.IndexType();
case 'bf16':
case 'f16':
case 'f32':
case 'f64':
case 'f80':
case 'f128':
case 'tf32':
case 'f8E5M2':
case 'f8E4M3':
case 'f8E4M3FN':
case 'f8E5M2FNUZ':
case 'f8E4M3FNUZ':
case 'f8E4M3B11FNUZ':
case 'f8E3M4':
case 'f8E8M0FNU':
case 'f4E2M1FN':
case 'f6E2M3FN':
case 'f6E3M2FN':
return new _.FloatType(value);
default:
this.emitError(`Invalid type '${value}'`);
}
break;
}
default: {
break;
}
}
this.emitError(`Invalid type '${this.getTokenSpelling().str()}'`);
return null;
}
parseExtendedType() {
return this.parseExtendedSymbol(this.state.asmState, this.state.typeAliasDefinitions, (dialectName, symbolData) => {
const dialect = this.context.getOrLoadDialect(dialectName);
if (dialect) {
const curLexerPos = this.getToken().loc.position;
this.resetToken(symbolData.data());
const customParser = new _.CustomDialectAsmParser(symbolData, this);
const type = dialect.parseType(customParser, dialectName);
this.resetToken(curLexerPos);
return type;
}
return new _.OpaqueType(dialectName, symbolData);
});
}
parseExtendedSymbol(asmState, aliases, createSymbol) {
const tok = this.getToken();
const startPos = tok.loc.position;
const tokSpelling = tok.getSpelling().str();
const isType = tokSpelling.startsWith('!');
const identifier = tok.getSpelling().drop_front();
this.consumeToken();
const dotIndex = identifier.indexOf('.');
let dialectName = identifier;
let symbolData = '';
if (dotIndex !== -1) {
dialectName = identifier.substring(0, dotIndex);
symbolData = identifier.substring(dotIndex + 1);
}
const isPrettyName = symbolData.length > 0 || identifier.endsWith('.');
const hasTrailingData = this.getToken().is(_.Token.less) && (startPos + tokSpelling.length) === this.getToken().loc.position;
if (!hasTrailingData && !isPrettyName) {
if (!aliases.has(tokSpelling)) {
this.emitError(`Undefined symbol alias '${identifier}'`);
}
if (asmState) {
if (isType) {
asmState.addTypeAliasUses(identifier);
} else {
asmState.addAttrAliasUses(identifier);
}
}
return aliases.get(tokSpelling);
}
if (hasTrailingData) {
if (this.getToken().is(_.Token.less)) {
symbolData += this.parseDialectSymbolBody();
}
if (!isPrettyName) {
symbolData = symbolData.slice(1, -1);
}
}
const posOffset = dotIndex === -1 ? dialectName.length : dotIndex;
symbolData = new _.StringRef(symbolData, startPos + 1 + posOffset + 1);
return createSymbol(dialectName, symbolData);
}
parseFunctionType() {
const inputs = this.parseTypeListParens();
this.parseToken(_.Token.arrow, "Expected '->' in function type");
const results = this.parseFunctionResultTypes();
return new _.FunctionType(inputs, results);
}
parseFunctionResultTypes() {
if (this.getToken().is(_.Token.l_paren)) {
return this.parseTypeListParens();
}
const type = this.parseNonFunctionType();
if (type) {
return [type];
}
return [];
}
parseCommaSeparatedList(delimiter, parseElement) {
const results = [];
const delimiters = {
none: [null, null],
paren: [_.Token.l_paren, _.Token.r_paren],
square: [_.Token.l_square, _.Token.r_square],
angle: [_.Token.less, _.Token.greater],
brace: [_.Token.l_brace, _.Token.r_brace],
optionalParen: [_.Token.l_paren, _.Token.r_paren],
optionalSquare: [_.Token.l_square, _.Token.r_square],
optionalAngle: [_.Token.less, _.Token.greater],
optionalBrace: [_.Token.l_brace, _.Token.r_brace]
};
const [open, close] = delimiters[delimiter] || [null, null];
const isOptional = delimiter && delimiter.startsWith('optional');
if (open) {
if (isOptional) {
if (!this.consumeIf(open)) {
return results;
}
} else {
this.parseToken(open, `Expected '${open}'`);
}
if (close && this.consumeIf(close)) {
return results;
}
}
const first = parseElement();
if (first !== null && first !== undefined) {
results.push(first);
}
while (this.consumeIf(_.Token.comma)) {
const elem = parseElement();
if (elem !== null && elem !== undefined) {
results.push(elem);
}
}
if (close) {
this.parseToken(close, `Expected '${close}'`);
}
return results;
}
parseFloatAttr(type = null, isNegative = false) {
const val = this.getToken().getFloatingPointValue();
if (val === null) {
this.emitError('Floating point value too large for attribute');
}
this.consumeToken(_.Token.floatliteral);
if (!type) {
if (this.consumeIf(_.Token.colon)) {
type = this.parseType();
} else {
type = new _.FloatType('f64');
}
}
if (type instanceof _.FloatType === false) {
this.emitError('Floating point value not valid for specified type');
}
return new _.FloatAttr(type, isNegative ? -val : val);
}
parseDecOrHexAttr(type = null, isNegative = false) {
const loc = this.getToken().loc.copy();
const tok = this.getToken();
const spelling = tok.getSpelling().str();
this.consumeToken(_.Token.integer);
if (!type) {
if (this.consumeIf(_.Token.colon)) {
type = this.parseType();
} else {
type = new _.IntegerType('i64');
}
}
if (type instanceof _.FloatType) {
return new _.FloatAttr(type, isNegative ? -spelling : spelling);
}
if (!(type instanceof _.IntegerType) && !(type instanceof _.IndexType)) {
this.emitError(loc, 'Integer literal not valid for specified type');
}
if (isNegative && type instanceof _.IntegerType && type.name.startsWith('ui')) {
this.emitError(loc, 'Negative integer literal not valid for unsigned integer type');
}
const val = tok.spelling.getAsInteger(spelling.length > 1 && spelling[1] === 'x' ? 0 : 10);
if (val === null) {
this.emitError(loc, 'Integer constant out of range for attribute');
}
return new _.IntegerAttr(type, isNegative ? -val : val);
}
parseAttribute(type = null) {
switch (this.getToken().kind) {
case _.Token.kw_affine_map: {
this.consumeToken(_.Token.kw_affine_map);
this.parseToken(_.Token.less, "Expected '<' in affine map");
const map = this.parseAffineMapReference();
this.parseToken(_.Token.greater, "Expected '>' in affine map");
return new _.AffineMapAttr(map);
}
case _.Token.kw_affine_set: {
this.consumeToken(_.Token.kw_affine_set);
this.parseToken(_.Token.less, "Expected '<' in integer set");
const set = this.parseIntegerSetReference();
this.parseToken(_.Token.greater, "Expected '>' in integer set");
return new _.IntegerSetAttr(set);
}
case _.Token.l_square: {
this.consumeToken(_.Token.l_square);
const elements = [];
this.parseCommaSeparatedListUntil(_.Token.r_square, () => {
elements.push(this.parseAttribute());
});
return new _.ArrayAttr(elements);
}
case _.Token.kw_false:
this.consumeToken(_.Token.kw_false);
return new _.BoolAttr(false);
case _.Token.kw_true:
this.consumeToken(_.Token.kw_true);
return new _.BoolAttr(true);
case _.Token.kw_dense:
return this.parseDenseElementsAttr(type);
case _.Token.kw_dense_resource:
return this.parseDenseResourceElementsAttr(type);
case _.Token.kw_array:
return this.parseDenseArrayAttr(type);
case _.Token.l_brace: {
const attributes = new Map();
this.parseAttributeDict(attributes);
return new _.DictionaryAttr(attributes);
}
case _.Token.hash_identifier:
return this.parseExtendedAttr(type);
case _.Token.floatliteral:
return this.parseFloatAttr(type, false);
case _.Token.integer:
return this.parseDecOrHexAttr(type, false);
case _.Token.minus: {
this.consumeToken(_.Token.minus);
if (this.getToken().is(_.Token.integer)) {
return this.parseDecOrHexAttr(type, true);
}
if (this.getToken().is(_.Token.floatliteral)) {
return this.parseFloatAttr(type, true);
}
this.emitError('Expected constant integer or floating point value');
break;
}
case _.Token.kw_loc:
return this.parseOptionalLocationSpecifier();
case _.Token.kw_sparse:
return this.parseSparseElementsAttr(type);
case _.Token.kw_strided:
return this.parseStridedLayoutAttr();
case _.Token.kw_distinct:
return this.parseDistinctAttr(type);
case _.Token.string: {
const value = this.getToken().getStringValue();
this.consumeToken(_.Token.string);
if (!type && this.consumeIf(_.Token.colon)) {
type = this.parseType();
}
return new _.StringAttr(value, type);
}
case _.Token.at_identifier: {
const nameStr = this.getToken().getSymbolReference();
this.consumeToken(_.Token.at_identifier);
const nestedRefs = [];
while (this.getToken().is(_.Token.colon)) {
const curPointer = this.getToken().loc.position;
this.consumeToken(_.Token.colon);
if (!this.consumeIf(_.Token.colon)) {
this.resetToken(curPointer);
break;
}
if (this.getToken().isNot(_.Token.at_identifier)) {
this.emitError('Expected nested symbol reference identifier');
}
nestedRefs.push(this.getToken().getSymbolReference());
this.consumeToken(_.Token.at_identifier);
}
return new _.SymbolRefAttr(nameStr, nestedRefs);
}
case _.Token.kw_unit:
this.consumeToken(_.Token.kw_unit);
return new _.UnitAttr();
case _.Token.bare_identifier: {
const tokenValue = this.getTokenSpelling().str();
if (tokenValue === 'tensor' || tokenValue === 'vector' || tokenValue === 'memref' ||
tokenValue === 'none' || tokenValue === 'index' || /^[su]?i[0-9]+$/.test(tokenValue) ||
/^f[0-9]+$/.test(tokenValue) || tokenValue === 'bf16' || tokenValue === 'tf32' ||
/^f\d+E\d+M\d+/.test(tokenValue) || tokenValue === 'complex' || tokenValue === 'tuple') {
const parsedType = this.parseType();
return { value: parsedType, type: new _.PrimitiveType('type') };
}
if (tokenValue === 'DEFAULT') {
this.consumeToken(_.Token.bare_identifier);
return { value: tokenValue };
}
this.consumeToken(_.Token.bare_identifier);
return { value: tokenValue };
}
case _.Token.exclamation_identifier: {
const parsedType = this.parseType();
return { value: parsedType, type: new _.PrimitiveType('type') };
}
case _.Token.percent_identifier: {
const value = this.getTokenSpelling().str();
this.consumeToken(_.Token.percent_identifier);
return { value };
}
case _.Token.less: {
const value = this.parseDialectSymbolBody();
return { value };
}
default: {
const parsedType = this.parseOptionalType();
if (parsedType) {
return new _.TypeAttrOf(parsedType);
}
this.emitError(`Unexpected attribute token '${this.getTokenSpelling().str()}'`);
break;
}
}
return null;
}
parseExtendedAttr(type = null) {
const attr = this.parseExtendedSymbol(this.state.asmState, this.state.attributeAliasDefinitions, (dialectName, symbolData) => {
let attrType = type;
if (this.consumeIf(_.Token.colon)) {
attrType = this.parseType();
if (!attrType) {
return new _.Attribute();
}
}
const dialect = this.context.getOrLoadDialect(dialectName);
if (dialect) {
const curLexerPos = this.getToken().loc.position;
this.resetToken(symbolData.data());
const customParser = new _.CustomDialectAsmParser(symbolData, this);
const attr = dialect.parseAttribute(customParser, attrType);
this.resetToken(curLexerPos);
if (attr) {
return attr;
}
}
return new _.OpaqueAttr(`#${dialectName}`, symbolData, attrType);
});
return attr;
}
parseResourceHandle(dialect) {
const name = this.parseOptionalKeywordOrString();
if (!name) {
this.emitError("Expected identifier key for 'resource' entry");
}
const resources = this.state.dialectResources;
if (!resources.has(dialect)) {
resources.set(dialect, new Map());
}
const dialectEntries = resources.get(dialect);
if (!dialectEntries.has(name)) {
const handle = dialect.declareResource(name);
const key = dialect.getResourceKey(handle);
dialectEntries.set(name, { key, handle });
}
const entry = dialectEntries.get(name);
return entry.handle;
}
parseDenseElementsAttr(attrType) {
this.consumeToken(_.Token.kw_dense);
this.parseToken(_.Token.less, "Expected '<' after 'dense'");
// Type-first syntax: dense<TYPE : [ATTR, ...]>
// Reference: AttributeParser.cpp parseDenseElementsAttrTyped()
if (!this.getToken().is(_.Token.l_paren)) {
const type = this.parseDenseElementsAttrTyped();
if (type) {
return type;
}
}
// Literal-first syntax: dense<[values]> : type
let literalParser = null;
if (!this.consumeIf(_.Token.greater)) {
literalParser = new _.TensorLiteralParser(this);
literalParser.parse(/* allowHex */ true);
this.parseToken(_.Token.greater, "Expected '>'");
}
const type = this.parseElementsLiteralType(attrType);
const value = literalParser ? literalParser.getAttr(type) : null;
return new _.DenseElementsAttr(value, type);
}
parseDenseElementsAttrTyped() {
const typeLoc = this.getToken().loc.copy();
const type = this.parseOptionalType();
if (!type) {
return null;
}
if (!this.consumeIf(_.Token.colon)) {
this.emitError("Expected ':' after type in dense attribute");
}
if (!(type instanceof _.RankedTensorType) && !(type instanceof _.VectorType)) {
this.emitError(typeLoc, 'Expected a shaped type for dense elements');
}
const values = [];
const parseElement = () => {
values.push(this.parseAttribute());
};
if (this.getToken().is(_.Token.l_square)) {
const parseElements = (shape) => {
if (shape.length <= 1) {
this.parseCommaSeparatedList('square', parseElement);
} else {
this.parseCommaSeparatedList('square', () => parseElements(shape.slice(1)));
}
};
parseElements(type.shape || []);
} else {
parseElement();
}
this.parseToken(_.Token.greater, "Expected '>' to close dense attribute");
const extractedValues = values.map((v) => {
if (v instanceof _.IntegerAttr) {
return typeof v.value === 'bigint' ? v.value : Number(v.value);
}
if (v instanceof _.FloatAttr) {
return parseFloat(v.value);
}
return v;
});
return new _.DenseElementsAttr(extractedValues, type);
}
parseDenseResourceElementsAttr(attrType) {
this.consumeToken(_.Token.kw_dense_resource);
this.parseToken(_.Token.less, "Expected '<' after 'dense_resource'");
const rawHandle = this.parseResourceHandle(this.context.getOrLoadDialect('builtin'));
this.parseToken(_.Token.greater, "Expected '>'");
let type = attrType;
if (!type) {
this.parseToken(_.Token.colon, "Expected ':'");
type = this.parseType();
}
return new _.DenseResourceElementsAttr(type, rawHandle);
}
parseDenseArrayAttr(/* attrType */) {
this.consumeToken(_.Token.kw_array);
this.parseToken(_.Token.less, "Expected '<' after 'array'");
const typeLoc = this.getToken().loc.copy();
const arrayType = this.parseType();
if (!(arrayType instanceof _.IntegerType) && !(arrayType instanceof _.FloatType)) {
this.emitError(typeLoc, `Expected integer or float type, got '${arrayType}'`);
}
if (this.consumeIf(_.Token.greater)) {
return new _.DenseArrayAttr(arrayType, 0, []);
}
this.parseToken(_.Token.colon, "Expected ':' after dense array type");
const arrayValues = [];
while (this.getToken().isNot(_.Token.greater)) {
const val = this.parseAttribute();
arrayValues.push(val && val.value !== undefined ? val.value : val);
this.consumeIf(_.Token.comma);
}
this.parseToken(_.Token.greater, "Expected '>' to close an array attribute");
return new _.DenseArrayAttr(arrayType, arrayValues.length, arrayValues);
}
parseSparseElementsAttr(attrType) {
this.consumeToken(_.Token.kw_sparse);
this.parseToken(_.Token.less, "Expected '<' after 'sparse'");
let indices = null;
let values = null;
if (!this.consumeIf(_.Token.greater)) {
const indiceParser = new _.TensorLiteralParser(this);
indiceParser.parse(/* allowHex */ false);
indices = indiceParser._storage.map((elem) => {
const val = elem.isNegative ? -elem.value : elem.value;
return typeof val === 'string' ? parseInt(val, 10) : val;
});
this.parseToken(_.Token.comma, "Expected ','");
const valuesParser = new _.TensorLiteralParser(this);
valuesParser.parse(/* allowHex */ true);
if (valuesParser._hexStorage) {
values = valuesParser._hexStorage;
} else {
values = valuesParser._storage.map((elem) => {
if (elem.kind === 'float') {
const val = parseFloat(elem.value);
return elem.isNegative ? -val : val;
}
const val = elem.isNegative ? -elem.value : elem.value;
return typeof val === 'string' ? parseFloat(val) : val;
});
}
this.parseToken(_.Token.greater, "Expected '>'");
}
const type = this.parseElementsLiteralType(attrType);
if (values && !(values instanceof Uint8Array) && type && type.elementType instanceof _.ComplexType) {
const packed = [];
for (let i = 0; i + 1 < values.length; i += 2) {
packed.push(new base.Complex(values[i], values[i + 1]));
}
values = packed;
}
return new _.SparseElementsAttr(type, indices, values);
}
parseStridedLayoutAttr() {
this.consumeToken(_.Token.kw_strided);
this.parseToken(_.Token.less, "Expected '<' after 'strided'");
this.parseToken(_.Token.l_square, "Expected '['");
// Parse dimension list: integer or '?' separated by commas
const strides = [];
while (this.getToken().isNot(_.Token.r_square)) {
const loc = this.getToken().loc.copy();
const value = this.parseOptionalInteger();
if (value !== null) {
strides.push(value);
} else if (this.consumeIf(_.Token.question)) {
strides.push('?');
} else {
this.emitError(loc, "Expected a 64-bit signed integer or '?'");
}
if (this.getToken().isNot(_.Token.r_square)) {
this.consumeIf(_.Token.comma);
}
}
this.parseToken(_.Token.r_square, "Expected ']'");
let offset = null;
if (this.consumeIf(_.Token.comma)) {
this.parseToken(_.Token.kw_offset, "Expected 'offset' after comma");
this.parseToken(_.Token.colon, "Expected ':' after 'offset'");
const offsetLoc = this.getToken().loc.copy();
offset = this.parseOptionalInteger();
if (offset === null) {
if (this.consumeIf(_.Token.question)) {
offset = '?';
} else {
this.emitError(offsetLoc, "Expected a 64-bit signed integer or '?' for offset in strided layout");
}
}
}
this.parseToken(_.Token.greater, "Expected '>'");
return new _.StridedLayoutAttr(offset, strides);
}
// Parse an affine map reference using AffineParser
// Following reference: Parser::parseAffineMapReference
parseAffineMapReference() {
const curLoc = this.getToken().loc.copy();
const affineParser = new _.AffineParser(this.state);
const result = affineParser.parseAffineMapOrIntegerSetInline();
if (result.set) {
this.emitError(curLoc, 'Expected AffineMap, but got IntegerSet');
}
return result.map;
}
// Parse an integer set reference using AffineParser
// Following reference: Parser::parseIntegerSetReference
parseIntegerSetReference() {
const curLoc = this.getToken().loc.copy();
const affineParser = new _.AffineParser(this.state);
const result = affineParser.parseAffineMapOrIntegerSetInline();
if (result.map) {
this.emitError(curLoc, 'Expected IntegerSet, but got AffineMap');
}
return result.set;
}
// Parse an AffineMap where dims/symbols are SSA ids
// Following reference: Parser::parseAffineMapOfSSAIds
parseAffineMapOfSSAIds(parseElement, delimiter = 'Paren') {
const affineParser = new _.AffineParser(this.state, true, parseElement);
return affineParser.parseAffineMapOfSSAIds(delimiter);
}
// Parse an AffineExpr where dims/symbols are SSA ids
// Following reference: Parser::parseAffineExprOfSSAIds
parseAffineExprOfSSAIds(parseElement) {
const affineParser = new _.AffineParser(this.state, true, parseElement);
return affineParser.parseAffineExprOfSSAIds();
}
parseDistinctAttr(type) {
this.consumeToken(_.Token.kw_distinct);
this.parseToken(_.Token.l_square, "Expected '[' after 'distinct'");
if (this.getToken().isNot(_.Token.integer)) {
this.emitError('Expected distinct ID');
}
const token = this.getToken();
const id = token.getUInt64IntegerValue();
if (id === null) {
this.emitError('Expected an unsigned 64-bit integer');
}
this.consumeToken(_.Token.integer);
this.parseToken(_.Token.r_square, "Expected ']' to close distinct ID");
this.parseToken(_.Token.less, "Expected '<' after distinct ID");
let referencedAttr = null;
if (this.consumeIf(_.Token.greater)) {
referencedAttr = new _.UnitAttr();
} else {
referencedAttr = this.parseAttribute(type);
this.parseToken(_.Token.greater, "Expected '>' to close distinct attribute");
}
return { value: `distinct[${id.toString()}]`, referencedAttr, type: 'distinct' };
}
parseElementsLiteralType(type) {
if (!type) {
this.parseToken(_.Token.colon, "Expected ':'");
return this.parseType();
}
// Type is a concrete type object - use it directly
return type;
}
parseOptionalAttribute(type) {
switch (this.getToken().kind) {
case _.Token.at_identifier:
case _.Token.floatliteral:
case _.Token.integer:
case _.Token.hash_identifier:
case _.Token.kw_affine_map:
case _.Token.kw_affine_set:
case _.Token.kw_dense:
case _.Token.kw_dense_resource:
case _.Token.kw_false:
case _.Token.kw_loc:
case _.Token.kw_sparse:
case _.Token.kw_true:
case _.Token.kw_unit:
case _.Token.l_brace:
case _.Token.l_square:
case _.Token.less:
case _.Token.minus:
case _.Token.string:
return this.parseAttribute(type);
default: {
const value = this.parseOptionalType(type);
if (value) {
return { value, type: 'type' };
}
return null;
}
}
}
parseOptionalInteger(type) {
// Parse `false` and `true` keywords as 0 and 1 respectively.
if (this.consumeIf(_.Token.kw_false)) {
return type === 'int64' ? 0n : 0;
}
if (this.consumeIf(_.Token.kw_true)) {
return type === 'int64' ? 1n : 1;
}
if (this.getToken().isNot(_.Token.integer) && this.getToken().isNot(_.Token.minus)) {
return null;
}
const negative = this.consumeIf(_.Token.minus);
const curTokLoc = this.getToken().loc.copy();
const curTok = this.getToken();
this.parseToken(_.Token.integer, 'Expected integer value');
if (type === 'int64') {
const str = curTok.spelling.str();
const result = BigInt(str);
return negative ? -result : result;
}
const spelling = curTok.spelling;
const isHex = spelling.length > 1 && spelling[1] === 'x';
const result = spelling.getAsInteger(isHex ? 0 : 10);
if (result === null) {
this.emitError(curTokLoc, 'Integer value too large');
}
return negative ? -result : result;
}
parseInteger(type) {
const result = this.parseOptionalInteger(type);
if (result === null) {
this.emitError('Expected integer value');
}
return result;
}
parseOptionalDecimalInteger() {
if (this.getToken().isNot(_.Token.integer) && this.getToken().isNot(_.Token.minus)) {
return null;
}
const negative = this.consumeIf(_.Token.minus);
const curTokLoc = this.getToken().loc.copy();
const curTok = this.getToken();
this.parseToken(_.Token.integer, 'Expected integer value');
const spelling = curTok.spelling.str();
if (spelling[0] === '0' && spelling.length > 1 && spelling[1] === 'x') {
this.state.lex.resetPointer(curTok.loc.position + 1);
this.state.curToken = this.state.lex.lexToken();
return 0;
}
const result = curTok.spelling.getAsInteger(10);
if (result === null) {
this.emitError(curTokLoc, 'Integer value too large');
}
return negative ? -result : result;
}
parseDecimalInteger() {
const result = this.parseOptionalDecimalInteger();
if (result === null) {
this.emitError('Expected decimal integer value');
}
return result;
}
parseOptionalVerticalBar() {
return this.consumeIf(_.Token.vertical_bar);
}
parseOptionalString() {
if (!this.getToken().is(_.Token.string)) {
return null;
}
const string = this.getToken().getStringValue();
this.consumeToken();
return string;
}
parseDialectSymbolBody() {
const decoder = this.state.lex.decoder;
const startPos = this.getToken().loc.position;
decoder.position = startPos;
const nestedPunctuation = [];
const chars = [];
do {
const c = decoder.decode();
if (c === undefined) {
if (nestedPunctuation.length > 0) {
this.emitError(`unbalanced '${nestedPunctuation[nestedPunctuation.length - 1]}' character in pretty dialect name`);
}
this.emitError('unexpected nul or EOF in pretty dialect name');
}
chars.push(c);
switch (c) {
case '<':
case '[':
case '(':
case '{':
nestedPunctuation.push(c);
continue;
case '-': {
const pos = decoder.position;
const next = decoder.decode();
if (next === '>') {
chars.push(next);
} else {
decoder.position = pos;
}
continue;
}
case '>':
if (nestedPunctuation[nestedPunctuation.length - 1] !== '<') {
this.emitError(`unbalanced '${nestedPunctuation[nestedPunctuation.length - 1]}' character in pretty dialect name`);
}
nestedPunctuation.pop();
break;
case ']':
if (nestedPunctuation[nestedPunctuation.length - 1] !== '[') {
this.emitError(`unbalanced '${nestedPunctuation[nestedPunctuation.length - 1]}' character in pretty dialect name`);
}
nestedPunctuation.pop();
break;
case ')':
if (nestedPunctuation[nestedPunctuation.length - 1] !== '(') {
this.emitError(`unbalanced '${nestedPunctuation[nestedPunctuation.length - 1]}' character in pretty dialect name`);
}
nestedPunctuation.pop();
break;
case '}':
if (nestedPunctuation[nestedPunctuation.length - 1] !== '{') {
this.emitError(`unbalanced '${nestedPunctuation[nestedPunctuation.length - 1]}' character in pretty dialect name`);
}
nestedPunctuation.pop();
break;
case '"': {
this.resetToken(decoder.position - 1);
if (this.getToken().isNot(_.Token.string)) {
this.emitError('expected string in dialect symbol body');
}
const spelling = this.getToken().getSpelling().str();
for (let i = 1; i < spelling.length; i++) {
chars.push(spelling[i]);
}
decoder.position = this.getToken().loc.position + spelling.length;
continue;
}
default:
continue;
}
} while (nestedPunctuation.length > 0);
this.resetToken(decoder.position);
return chars.join('');
}
getEncodedSourceLocation(loc) {
return this.state.lex.getEncodedSourceLocation(loc);
}
getToken() {
return this.state.curToken;
}
getTokenSpelling() {
return this.state.curToken.spelling;
}
resetToken(tokPos) {
this.state.lex.resetPointer(tokPos);
this.state.curToken = this.state.lex.lexToken();
}
consumeToken(kind) {
if (kind !== undefined && this.state.curToken.kind !== kind) {
this.emitError(`consumeToken: Expected '${kind}'`);
}
this.state.curToken = this.state.lex.lexToken();
}
parseToken(kind, message) {
if (this.consumeIf(kind)) {
return;
}
this.emitError(message || `Expected '${kind}'`);
}
consumeIf(kind) {
if (this.state.curToken.isNot(kind)) {
return false;
}
this.consumeToken(kind);
return true;
}
emitError(...args) {
const loc = args.length === 2 ? args[0] : this.getToken().loc;
const message = args.length === 2 ? args[1] : args[0];
const location = this.getEncodedSourceLocation(loc);
throw new mlir.Error(`${message} at ${location.impl.printLocationInternal()}.`);
}
emitWarning() {
}
static parseSymbol(inputStr, context, parserFn) {
const decoder = text.Decoder.open(inputStr);
const config = new _.ParserConfig(context);
const sourceMgr = new _.SourceMgr();
sourceMgr.addNewSourceBuffer(new _.MemoryBuffer(decoder, ''));
const state = new _.ParserState(sourceMgr, config);
const parser = new _.Parser(state, context);
const startPos = state.curToken.position || 0;
const symbol = parserFn(parser);
if (!symbol) {
return { symbol: null, numRead: 0 };
}
const endPos = parser.getToken().loc.position;
const numRead = endPos - startPos;
return { symbol, numRead };
}
static parseAttribute(attrStr, context, type = null) {
const result = _.Parser.parseSymbol(attrStr, context, (parser) => {
return parser.parseAttribute(type);
});
if (!result.symbol) {
return { attribute: null, numRead: 0 };
}
return { attribute: result.symbol, numRead: result.numRead };
}
static parseType(typeStr, context) {
const result = _.Parser.parseSymbol(typeStr, context, (parser) => {
return parser.parseType();
});
if (!result.symbol) {
return { type: null, numRead: 0 };
}
return { type: result.symbol, numRead: result.numRead };
}
};
_.TopLevelOperationParser = class extends _.Parser {
parse(block) {
const opParser = new _.OperationParser(this.state);
while (true) {
switch (this.getToken().kind) {
case _.Token.eof: {
opParser.finalize(block);
return block;
}
case _.Token.hash_identifier: {
this.parseAttributeAliasDef();
break;
}
case _.Token.exclamation_identifier: {
this.parseTypeAliasDef();
break;
}
case _.Token.file_metadata_begin: {
this.parseFileMetadataDictionary();
break;
}
default: {
const op = opParser.parseOperation();
block.operations.push(op);
}
}
}
}
parseAttributeAliasDef() {
const aliasName = this.getToken().getSpelling().str();
this.consumeToken(_.Token.hash_identifier);
this.parseToken(_.Token.equal, "Expected '=' in attribute alias definition");
let attr = null;
if (this.getToken().is(_.Token.l_paren)) {
const map = this.parseAffineMapReference();
attr = new _.AffineMapAttr(map);
} else {
attr = this.parseAttribute();
}
this.state.attributeAliasDefinitions.set(aliasName, attr);
}
parseTypeAliasDef() {
const aliasName = this.getTokenSpelling().str();
this.consumeToken(_.Token.exclamation_identifier);
this.parseToken(_.Token.equal, "Expected '=' in type alias definition");
// Legacy MLIR syntax: !alias = type <actual-type>
if (this.getToken().is(_.Token.bare_identifier) && this.getTokenSpelling().str() === 'type') {
this.consumeToken(_.Token.bare_identifier);
}
const type = this.parseType();
this.state.typeAliasDefinitions.set(aliasName, type);
}
parseFileMetadataDictionary() {
this.consumeToken(_.Token.file_metadata_begin);
this.parseCommaSeparatedListUntil(_.Token.file_metadata_end, () => {
const keyLoc = this.getToken().loc;
const key = this.parseOptionalKeyword();
if (!key) {
this.emitError('Expected identifier key in file metadata dictionary');
}
this.parseToken(_.Token.colon, "Expected ':'");
if (key === 'dialect_resources') {
this.parseDialectResourceFileMetadata();
} else if (key === 'external_resources') {
this.parseExternalResourceFileMetadata();
} else {
this.emitError(keyLoc, `Unknown key '${key}' in file metadata dictionary`);
}
});
}
};
// Specialized parser for affine structures (affine maps, affine expressions, integer sets)
// Following the reference implementation in AffineParser.cpp
_.AffineParser = class extends _.Parser {
constructor(state, allowParsingSSAIds = false, parseElement = null) {
super(state);
this.allowParsingSSAIds = allowParsingSSAIds;
this.parseElement = parseElement;
this.dimsAndSymbols = []; // Array of {name, expr} pairs
this.numDimOperands = 0;
this.numSymbolOperands = 0;
}
// Parse an ambiguous affine map or integer set inline
// affine-map-or-integer-set ::= dim-and-symbol-id-lists (`->` | `:`) ...
parseAffineMapOrIntegerSetInline() {
const { numDims, numSymbols } = this.parseDimAndOptionalSymbolIdList();
if (this.consumeIf(_.Token.arrow)) {
return { map: this.parseAffineMapRange(numDims, numSymbols), set: null };
}
if (!this.consumeIf(_.Token.colon)) {
this.emitError("Expected '->' or ':'");
}
return { map: null, set: this.parseIntegerSetConstraints(numDims, numSymbols) };
}
// Parse dimension and optional symbol identifier lists: (d0, d1, ...)[s0, s1, ...]
parseDimAndOptionalSymbolIdList() {
const numDims = this.parseDimIdList();
let numSymbols = 0;
if (this.getToken().is(_.Token.l_square)) {
numSymbols = this.parseSymbolIdList();
}
return { numDims, numSymbols };
}
// Parse dimension identifier list: (d0, d1, ...)
parseDimIdList() {
let numDims = 0;
this.parseToken(_.Token.l_paren, "Expected '(' in dimensional identifier list");
if (this.getToken().isNot(_.Token.r_paren)) {
do {
const dimExpr = this.getAffineDimExpr(numDims);
this.parseIdentifierDefinition(dimExpr);
numDims++;
} while (this.consumeIf(_.Token.comma));
}
this.parseToken(_.Token.r_paren, "Expected ')' in dimensional identifier list");
return numDims;
}
// Parse symbol identifier list: [s0, s1, ...]
parseSymbolIdList() {
let numSymbols = 0;
this.parseToken(_.Token.l_square, "Expected '[' in symbol list");
if (this.getToken().isNot(_.Token.r_square)) {
do {
const symbolExpr = this.getAffineSymbolExpr(numSymbols);
this.parseIdentifierDefinition(symbolExpr);
numSymbols++;
} while (this.consumeIf(_.Token.comma));
}
this.parseToken(_.Token.r_square, "Expected ']' in symbol list");
return numSymbols;
}
parseIdentifierDefinition(idExpr) {
let name = null;
if (this.getToken().is(_.Token.bare_identifier)) {
name = this.getTokenSpelling().str();
this.consumeToken(_.Token.bare_identifier);
} else if (this.getToken().is(_.Token.inttype)) {
name = this.getTokenSpelling().str();
this.consumeToken(_.Token.inttype);
} else if (this.getToken().isKeyword()) {
name = this.getToken().getSpelling().str();
this.consumeToken();
} else {
this.emitError('Expected bare identifier');
}
for (const entry of this.dimsAndSymbols) {
if (entry.name === name) {
this.emitError(`Redefinition of identifier '${name}'`);
}
}
this.dimsAndSymbols.push({ name, expr: idExpr });
}
parseAffineMapRange(numDims, numSymbols) {
const exprs = [];
this.parseToken(_.Token.l_paren, "Expected '('");
if (this.getToken().isNot(_.Token.r_paren)) {
do {
const expr = this.parseAffineExpr();
if (!expr) {
this.emitError('Failed to parse affine expression');
}
exprs.push(expr);
} while (this.consumeIf(_.Token.comma));
}
this.parseToken(_.Token.r_paren, "Expected ')'");
return _.AffineMap.get(numDims, numSymbols, exprs);
}
parseIntegerSetConstraints(numDims, numSymbols) {
const constraints = [];
const isEqs = [];
this.parseToken(_.Token.l_paren, "Expected '('");
if (this.getToken().isNot(_.Token.r_paren)) {
do {
const { expr, isEq } = this.parseAffineConstraint();
constraints.push(expr);
isEqs.push(isEq);
} while (this.consumeIf(_.Token.comma));
}
this.parseToken(_.Token.r_paren, "Expected ')'");
// If no constraints, return degenerate true set (0 == 0)
if (constraints.length === 0) {
const zero = this.getAffineConstantExpr(0);
return _.IntegerSet.get(numDims, numSymbols, [zero], [true]);
}
return _.IntegerSet.get(numDims, numSymbols, constraints, isEqs);
}
parseAffineConstraint() {
const lhsExpr = this.parseAffineExpr();
if (!lhsExpr) {
this.emitError('Expected affine expression');
}
if (this.consumeIf(_.Token.greater) && this.consumeIf(_.Token.equal)) {
const rhsExpr = this.parseAffineExpr();
return { expr: this.makeSubExpr(lhsExpr, rhsExpr), isEq: false };
}
if (this.consumeIf(_.Token.less) && this.consumeIf(_.Token.equal)) {
const rhsExpr = this.parseAffineExpr();
return { expr: this.makeSubExpr(rhsExpr, lhsExpr), isEq: false };
}
if (this.consumeIf(_.Token.equal) && this.consumeIf(_.Token.equal)) {
const rhsExpr = this.parseAffineExpr();
return { expr: this.makeSubExpr(lhsExpr, rhsExpr), isEq: true };
}
this.emitError("Expected '>=', '<=', or '==' after affine expression");
return null;
}
makeSubExpr(lhs, rhs) {
const negOne = this.getAffineConstantExpr(-1);
const negRhs = this.getAffineBinaryOpExpr(_.AffineExprKind.Mul, negOne, rhs);
return this.getAffineBinaryOpExpr(_.AffineExprKind.Add, lhs, negRhs);
}
parseAffineExpr() {
return this.parseAffineLowPrecOpExpr(null, null);
}
parseAffineLowPrecOpExpr(llhs, llhsOp) {
const lhs = this.parseAffineOperandExpr(llhs);
if (!lhs) {
return null;
}
const lOp = this.consumeIfLowPrecOp();
if (lOp) {
if (llhs) {
const sum = this.makeAffineBinaryOp(llhsOp, llhs, lhs);
return this.parseAffineLowPrecOpExpr(sum, lOp);
}
return this.parseAffineLowPrecOpExpr(lhs, lOp);
}
const hOp = this.consumeIfHighPrecOp();
if (hOp) {
const highRes = this.parseAffineHighPrecOpExpr(lhs, hOp);
if (!highRes) {
return null;
}
const expr = llhs ? this.makeAffineBinaryOp(llhsOp, llhs, highRes) : highRes;
const nextOp = this.consumeIfLowPrecOp();
if (nextOp) {
return this.parseAffineLowPrecOpExpr(expr, nextOp);
}
return expr;
}
if (llhs) {
return this.makeAffineBinaryOp(llhsOp, llhs, lhs);
}
return lhs;
}
parseAffineHighPrecOpExpr(llhs, llhsOp) {
const lhs = this.parseAffineOperandExpr(llhs);
if (!lhs) {
return null;
}
const op = this.consumeIfHighPrecOp();
if (op) {
if (llhs) {
const expr = this.makeAffineBinaryOp(llhsOp, llhs, lhs);
return this.parseAffineHighPrecOpExpr(expr, op);
}
return this.parseAffineHighPrecOpExpr(lhs, op);
}
if (llhs) {
return this.makeAffineBinaryOp(llhsOp, llhs, lhs);
}
return lhs;
}
consumeIfLowPrecOp() {
switch (this.getToken().kind) {
case _.Token.plus:
this.consumeToken(_.Token.plus);
return _.AffineLowPrecOp.Add;
case _.Token.minus:
this.consumeToken(_.Token.minus);
return _.AffineLowPrecOp.Sub;
default:
return _.AffineLowPrecOp.LNoOp;
}
}
consumeIfHighPrecOp() {
switch (this.getToken().kind) {
case _.Token.star:
this.consumeToken(_.Token.star);
return _.AffineHighPrecOp.Mul;
case _.Token.kw_floordiv:
this.consumeToken(_.Token.kw_floordiv);
return _.AffineHighPrecOp.FloorDiv;
case _.Token.kw_ceildiv:
this.consumeToken(_.Token.kw_ceildiv);
return _.AffineHighPrecOp.CeilDiv;
case _.Token.kw_mod:
this.consumeToken(_.Token.kw_mod);
return _.AffineHighPrecOp.Mod;
default:
return _.AffineHighPrecOp.HNoOp;
}
}
// Create affine binary op expression
makeAffineBinaryOp(op, lhs, rhs) {
if (op === 'Sub') {
// Subtraction: lhs - rhs = lhs + (-1 * rhs)
const negOne = this.getAffineConstantExpr(-1);
const negRhs = this.getAffineBinaryOpExpr(_.AffineExprKind.Mul, negOne, rhs);
return this.getAffineBinaryOpExpr(_.AffineExprKind.Add, lhs, negRhs);
}
return this.getAffineBinaryOpExpr(op, lhs, rhs);
}
parseAffineOperandExpr(lhs) {
if (this.getToken().is(_.Token.kw_symbol)) {
return this.parseSymbolSSAIdExpr();
}
if (this.getToken().is(_.Token.percent_identifier)) {
return this.parseSSAIdExpr(false);
}
if (this.getToken().is(_.Token.integer)) {
return this.parseIntegerExpr();
}
if (this.getToken().is(_.Token.l_paren)) {
return this.parseParentheticalExpr();
}
if (this.getToken().is(_.Token.minus)) {
return this.parseNegateExpression(lhs);
}
if (this.getToken().isAny(_.Token.bare_identifier, _.Token.inttype)) {
return this.parseBareIdExpr();
}
if (this.getToken().isKeyword()) {
return this.parseBareIdExpr();
}
if (this.getToken().isAny(_.Token.plus, _.Token.star)) {
if (lhs) {
this.emitError('Missing right operand of binary operator');
}
this.emitError('Missing left operand of binary operator');
}
if (lhs) {
this.emitError('Missing right operand of binary operator');
}
this.emitError('Expected affine expression');
return null;
}
parseSymbolSSAIdExpr() {
this.parseToken(_.Token.kw_symbol);
this.parseToken(_.Token.l_paren);
const symbolExpr = this.parseSSAIdExpr(true);
this.parseToken(_.Token.r_paren);
return symbolExpr;
}
parseSSAIdExpr(isSymbol) {
if (!this.allowParsingSSAIds) {
this.emitError('Unexpected SSA identifier');
}
if (this.getToken().isNot(_.Token.percent_identifier)) {
this.emitError('Expected SSA identifier');
}
const name = this.getTokenSpelling().str();
for (const entry of this.dimsAndSymbols) {
if (entry.name === name) {
this.consumeToken(_.Token.percent_identifier);
return entry.expr;
}
}
if (this.parseElement) {
this.parseElement(isSymbol);
} else {
this.consumeToken(_.Token.percent_identifier);
}
const idExpr = isSymbol
? this.getAffineSymbolExpr(this.numSymbolOperands++)
: this.getAffineDimExpr(this.numDimOperands++);
this.dimsAndSymbols.push({ name, expr: idExpr });
return idExpr;
}
parseIntegerExpr() {
const val = this.getToken().getUInt64IntegerValue();
this.consumeToken(_.Token.integer);
return this.getAffineConstantExpr(Number(val));
}
parseParentheticalExpr() {
this.parseToken(_.Token.l_paren, "Expected '('");
if (this.getToken().is(_.Token.r_paren)) {
this.emitError('No expression inside parentheses');
}
const expr = this.parseAffineExpr();
this.parseToken(_.Token.r_paren, "Expected ')'");
return expr;
}
parseNegateExpression(lhs) {
this.parseToken(_.Token.minus, "Expected '-'");
const operand = this.parseAffineOperandExpr(lhs);
if (!operand) {
this.emitError('Missing operand of negation');
}
const negOne = this.getAffineConstantExpr(-1);
return this.getAffineBinaryOpExpr(_.AffineExprKind.Mul, negOne, operand);
}
parseBareIdExpr() {
if (!this.isCurrentTokenAKeyword()) {
this.emitError('Expected bare identifier');
}
const name = this.getTokenSpelling().str();
this.consumeToken();
for (const entry of this.dimsAndSymbols) {
if (entry.name === name) {
return entry.expr;
}
}
this.emitError(`Use of undeclared identifier '${name}'`);
return null;
}
parseAffineMapOfSSAIds(delimiter = 'Paren') {
const exprs = [];
const open = delimiter === 'Paren' ? _.Token.l_paren : _.Token.l_square;
const close = delimiter === 'Paren' ? _.Token.r_paren : _.Token.r_square;
this.parseToken(open, `Expected '${open === _.Token.l_paren ? '(' : '['}'`);
if (this.getToken().isNot(close)) {
do {
const expr = this.parseAffineExpr();
exprs.push(expr);
} while (this.consumeIf(_.Token.comma));
}
this.parseToken(close, `Expected '${close === _.Token.r_paren ? ')' : ']'}'`);
return _.AffineMap.get(this.numDimOperands, this.dimsAndSymbols.length - this.numDimOperands, exprs);
}
parseAffineExprOfSSAIds() {
return this.parseAffineExpr();
}
getAffineDimExpr(position) {
return new _.AffineDimExpr(position);
}
getAffineSymbolExpr(position) {
return new _.AffineSymbolExpr(position);
}
getAffineConstantExpr(constant) {
return new _.AffineConstantExpr(constant);
}
getAffineBinaryOpExpr(kind, lhs, rhs) {
return new _.AffineBinaryOpExpr(kind, lhs, rhs);
}
};
_.OperationParser = class extends _.Parser {
constructor(state) {
super(state);
this.isolatedNameScopes = [];
this.pushSSANameScope(/* isIsolated */ true);
this._redirect = new Map([
['builtin.func', 'func.func'],
['builtin.constant', 'arith.constant'],
['builtin.return', 'func.return'],
['builtin.select', 'arith.select'],
['scf.select', 'arith.select'],
['scf.call', 'func.call'],
['builtin.view', 'memref.view'],
['builtin.dealloc', 'memref.dealloc'], ['func.dealloc', 'memref.dealloc'],
// Arith operations (from both builtin and func default dialects)
['builtin.addi', 'arith.addi'], ['func.addi', 'arith.addi'],
['builtin.subi', 'arith.subi'], ['func.subi', 'arith.subi'],
['builtin.muli', 'arith.muli'], ['func.muli', 'arith.muli'],
['builtin.divi_signed', 'arith.divsi'], ['func.divi_signed', 'arith.divsi'],
['builtin.divi_unsigned', 'arith.divui'], ['func.divi_unsigned', 'arith.divui'],
['builtin.divsi', 'arith.divsi'], ['func.divsi', 'arith.divsi'],
['builtin.divui', 'arith.divui'], ['func.divui', 'arith.divui'],
['builtin.remi_signed', 'arith.remsi'], ['func.remi_signed', 'arith.remsi'],
['builtin.remi_unsigned', 'arith.remui'], ['func.remi_unsigned', 'arith.remui'],
['builtin.andi', 'arith.andi'], ['func.andi', 'arith.andi'],
['builtin.ori', 'arith.ori'], ['func.ori', 'arith.ori'],
['builtin.xori', 'arith.xori'], ['func.xori', 'arith.xori'],
['builtin.shli', 'arith.shli'], ['func.shli', 'arith.shli'],
['builtin.shrsi', 'arith.shrsi'], ['func.shrsi', 'arith.shrsi'],
['builtin.shrui', 'arith.shrui'], ['func.shrui', 'arith.shrui'],
['builtin.addf', 'arith.addf'], ['func.addf', 'arith.addf'],
['builtin.subf', 'arith.subf'], ['func.subf', 'arith.subf'],
['builtin.mulf', 'arith.mulf'], ['func.mulf', 'arith.mulf'],
['builtin.divf', 'arith.divf'], ['func.divf', 'arith.divf'],
['builtin.cmpi', 'arith.cmpi'], ['func.cmpi', 'arith.cmpi'],
['builtin.cmpf', 'arith.cmpf'], ['func.cmpf', 'arith.cmpf'],
['builtin.index_cast', 'arith.index_cast'], ['func.index_cast', 'arith.index_cast'],
['builtin.sitofp', 'arith.sitofp'], ['func.sitofp', 'arith.sitofp'],
['builtin.fptosi', 'arith.fptosi'], ['func.fptosi', 'arith.fptosi'],
['builtin.truncf', 'arith.truncf'], ['func.truncf', 'arith.truncf'],
['builtin.extf', 'arith.extf'], ['func.extf', 'arith.extf'],
['builtin.splat', 'vector.splat'],
['func.splat', 'vector.splat'],
['scf.splat', 'vector.splat'],
// Memref operations
['builtin.alloc', 'memref.alloc'], ['func.alloc', 'memref.alloc'],
['builtin.load', 'memref.load'], ['func.load', 'memref.load'],
['builtin.store', 'memref.store'], ['func.store', 'memref.store'],
['builtin.subview', 'memref.subview'], ['func.subview', 'memref.subview'],
['builtin.dim', 'memref.dim'], ['func.dim', 'memref.dim'],
['builtin.view', 'memref.view'], ['func.view', 'memref.view'],
// Control flow operations
['builtin.cond_br', 'cf.cond_br'], ['func.cond_br', 'cf.cond_br'],
['builtin.br', 'cf.br'], ['func.br', 'cf.br'],
['builtin.switch', 'cf.switch'], ['func.switch', 'cf.switch'],
['builtin.assert', 'cf.assert'], ['func.assert', 'cf.assert'],
// AMX -> X86 AMX redirects
['amx.tile_load', 'x86.amx.tile_load'],
['amx.tile_mulf', 'x86.amx.tile_mulf'],
['amx.tile_muli', 'x86.amx.tile_muli'],
['amx.tile_store', 'x86.amx.tile_store'],
['amx.tile_zero', 'x86.amx.tile_zero'],
// Other redirects
['flow.constant', 'flow.tensor.constant'],
['util.initializer.return', 'util.return']
]);
}
finalize(block) {
const attributeAliases = this.state.attributeAliasDefinitions;
const deferredRefs = this.state.deferredLocsReferences;
const resolveLocation = (opOrArgument) => {
const fwdLoc = opOrArgument.loc;
if (fwdLoc instanceof _.OpaqueLoc) {
const locInfo = deferredRefs[fwdLoc.index];
const identifier = `#${locInfo.identifier}`;
const attr = attributeAliases.get(identifier);
if (!attr) {
throw new mlir.Error(`Operation location alias '${locInfo.identifier}' was never defined.`);
}
if (attr instanceof _.LocationAttr === false) {
throw new mlir.Error(`Expected location but found '${attr}'.`);
}
opOrArgument.loc = attr;
}
};
// Walk all operations and resolve locations on ops and block arguments
const walk = (operations) => {
for (const op of operations) {
if (op.loc) {
resolveLocation(op);
}
if (op.regions) {
for (const region of op.regions) {
if (!region.blocks) {
continue;
}
for (const blk of region.blocks) {
if (blk.arguments) {
for (const arg of blk.arguments) {
if (arg.loc) {
resolveLocation(arg);
}
}
}
if (blk.operations) {
walk(blk.operations);
}
}
}
}
}
};
walk(block.operations);
}
parseOperation() {
const loc = this.getToken().loc.copy();
const resultIDs = [];
let numExpectedResults = 0;
if (this.getToken().is(_.Token.percent_identifier)) {
const parseNextResult = () => {
const name = this.getTokenSpelling().str();
this.consumeToken(_.Token.percent_identifier);
let expectedSubResults = 1;
if (this.consumeIf(_.Token.colon)) {
if (this.getToken().isNot(_.Token.integer)) {
this.emitError('Expected integer number of results');
}
const val = parseInt(this.getToken().getSpelling().str(), 10);
if (!Number.isInteger(val) || val < 1) {
this.emitError('Expected named operation to have at least 1 result');
}
this.consumeToken(_.Token.integer);
expectedSubResults *= val;
}
resultIDs.push([name, expectedSubResults, null]);
numExpectedResults += expectedSubResults;
return true;
};
this.parseCommaSeparatedList('none', parseNextResult);
this.parseToken(_.Token.equal, "Expected '=' after SSA name");
}
let op = null;
// Reference Parser.cpp:1305: nameTok.is(Token::bare_identifier) || nameTok.isKeyword()
const nameTok = this.getToken();
if (nameTok.is(_.Token.bare_identifier) || nameTok.isKeyword()) {
op = this.parseCustomOperation(resultIDs);
} else if (nameTok.is(_.Token.string)) {
op = this.parseGenericOperation();
} else {
this.emitError(this.getToken().is(_.Token.eof) ? 'Unexpected end of input' : `Unexpected operation name '${this.getTokenSpelling().str()}'`);
}
if (!op) {
this.emitError('Failed to parse operation');
}
if (resultIDs.length > 0 && op.results.length === 0) {
this.emitError(loc, 'Cannot name an operation with no results');
}
if (resultIDs.length > 0 && op.results.length !== numExpectedResults) {
this.emitError(loc, `Operation '${op.name.getStringRef()}' defines '${op.results.length}' results but was provided '${numExpectedResults}' to bind`);
}
let index = 0;
for (const resIt of resultIDs) {
for (let subRes = 0; subRes < resIt[1]; subRes++) {
const result = op.results[index++];
if (result.name) {
this.emitError(`Result '${result.name}' already has name`);
}
// Visualization-specific addition to store the name for display
result.name = subRes === 0 ? resIt[0] : `${resIt[0]}.${subRes}`;
this.addDefinition(new _.UnresolvedOperand(resIt[2], resIt[0], subRes), result);
}
}
return op;
}
parseCustomOperation(resultIDs) {
const opLoc = this.getToken().loc.copy();
const opNameInfo = this.parseCustomOperationName();
const srcLocation = this.getEncodedSourceLocation(opLoc).impl;
const opState = new _.OperationState(srcLocation, opNameInfo);
delete opNameInfo.identifier; // Visualization-specific addition to keep the original name
const defaultDialect = (opNameInfo && opNameInfo.metadata && opNameInfo.metadata.defaultDialect) || this.state.defaultDialectStack[this.state.defaultDialectStack.length - 1];
this.state.defaultDialectStack.push(defaultDialect);
const customParser = new _.CustomOpAsmParser(resultIDs, this);
if (!opNameInfo.dialect.parseOperation(customParser, opState)) {
this.state.defaultDialectStack.pop();
this.emitError(opLoc, `Unsupported custom operation '${opState.identifier}'`);
}
if (!opNameInfo.metadata.hasParser && !opNameInfo.metadata.hasCustomAssemblyFormat && opNameInfo.metadata.assemblyFormat && opState.compatibility !== true) {
this.emitError(`Operation '${opState.identifier}' has assembly format but was handled by custom dialect code.`);
}
const trailingLoc = this.parseTrailingLocationSpecifier();
if (trailingLoc) {
opState.location = trailingLoc;
}
this.state.defaultDialectStack.pop();
return _.Operation.create(opState);
}
parseCustomOperationName() {
const nameTok = this.getToken();
const nameLoc = nameTok.loc.copy();
if (nameTok.kind !== _.Token.bare_identifier && !nameTok.isKeyword()) {
this.emitError('Expected bare identifier or keyword');
}
const identifier = nameTok.getSpelling().str(); // Visualization-specific addition to keep the original name
this.consumeToken();
let opName = this.redirect(identifier);
let opInfo = _.RegisteredOperationName.lookup(opName, this.context);
if (opInfo) {
opInfo.identifier = identifier;
return opInfo;
}
let index = opName.indexOf('.');
if (index === -1) {
const dialect = this.state.defaultDialectStack[this.state.defaultDialectStack.length - 1];
opName = this.redirect(`${dialect}.${opName}`);
}
index = opName.indexOf('.');
if (index === -1) {
this.emitError(nameLoc, `No dialect found '${opName}'`);
}
const dialectName = opName.substring(0, index);
const dialect = this.context.getOrLoadDialect(dialectName);
if (!dialect) {
if (this.context.undocumented(dialectName)) {
this.emitError(nameLoc, `Undocumented dialect '${dialectName}'`);
} else {
this.emitError(nameLoc, `Unsupported dialect '${dialectName}'`);
}
}
// Normalize operation name to canonical dialect name for metadata lookup
// (e.g., spv.Load -> spirv.Load when dialect.name is spirv)
if (dialectName !== dialect.name) {
opName = opName.replace(`${dialectName}.`, `${dialect.name}.`);
}
opInfo = _.RegisteredOperationName.lookup(opName, this.context);
if (!opInfo) {
this.emitError(nameLoc, `Unsupported operation '${opName}'`);
}
opInfo.identifier = identifier; // Visualization-specific addition to keep the original name
return opInfo;
}
parseGenericOperation() {
const opLoc = this.getToken().loc.copy();
const identifier = this.getToken().getStringValue();
if (identifier.length === 0) {
this.emitError('Empty operation name is invalid');
}
if (identifier.indexOf('\0') !== -1) {
this.emitError('Null character not allowed in operation name');
}
this.consumeToken(_.Token.string);
const name = this.redirect(identifier);
let opName = _.RegisteredOperationName.lookup(name, this.context);
if (!opName) {
const index = name.indexOf('.');
const dialect = index === -1 ? null : this.context.getOrLoadDialect(name.substring(0, index));
opName = new _.OperationName(dialect || null, name);
}
opName.identifier = identifier; // Visualization-specific addition to keep the original name
const srcLocation = this.getEncodedSourceLocation(opLoc).impl;
const result = new _.OperationState(srcLocation, opName);
this.parseGenericOperationAfterOpName(result);
return _.Operation.create(result);
}
parseGenericOperationAfterOpName(result) {
this.parseToken(_.Token.l_paren, "Expected '(' to start operand list");
const unresolvedOperands = [];
this.parseOptionalSSAUseList(unresolvedOperands);
this.parseToken(_.Token.r_paren, "Expected ')' to end operand list");
if (this.getToken().is(_.Token.l_square)) {
result.successors = [];
this.parseSuccessors(result.successors);
}
if (this.consumeIf(_.Token.less)) {
result.propertiesAttr = this.parseAttribute();
if (!result.propertiesAttr) {
this.emitError('Expected attribute as properties');
}
this.parseToken(_.Token.greater, "Expected '>' to close properties");
}
if (this.consumeIf(_.Token.l_paren)) {
do {
const region = result.addRegion();
this.parseRegion(region, undefined, false);
} while (this.consumeIf(_.Token.comma));
this.parseToken(_.Token.r_paren, "Expected ')' to end region list");
}
if (this.getToken().is(_.Token.l_brace)) {
this.parseAttributeDict(result.attributes);
}
this.parseToken(_.Token.colon, "Expected ':' followed by operation type");
const typeLoc = this.getToken().loc.copy();
const fnType = this.parseType();
if (fnType instanceof _.FunctionType === false) {
this.emitError(typeLoc, 'Expected function type');
}
const operandTypes = fnType.inputs;
if (operandTypes.length !== unresolvedOperands.length) {
this.emitError(typeLoc, `Expected ${unresolvedOperands.length} operand type${unresolvedOperands.length === 1 ? '' : 's'} but had ${operandTypes.length}`);
}
for (let i = 0; i < unresolvedOperands.length; i++) {
const unresolvedOperand = unresolvedOperands[i];
const type = operandTypes[i] || null;
const value = this.resolveSSAUse(unresolvedOperand, type);
result.operands.push(value);
}
result.addTypes(fnType.results);
const trailingLoc = this.parseTrailingLocationSpecifier();
if (trailingLoc) {
result.location = trailingLoc;
}
}
pushSSANameScope(isIsolated) {
if (isIsolated) {
this.isolatedNameScopes.push(new _.IsolatedSSANameScope());
}
this.isolatedNameScopes[this.isolatedNameScopes.length - 1].pushSSANameScope();
}
popSSANameScope() {
const currentNameScope = this.isolatedNameScopes[this.isolatedNameScopes.length - 1];
if (currentNameScope.definitionsPerScope.length === 1) {
this.isolatedNameScopes.pop();
} else {
currentNameScope.popSSANameScope();
}
}
addDefinition(useInfo, value) {
const entries = this.getSSAValueEntry(useInfo.name);
if (entries.length <= useInfo.number) {
entries.length = useInfo.number + 1;
}
entries[useInfo.number] = { value, loc: useInfo.location };
this.recordDefinition(useInfo.name);
}
recordDefinition(def) {
this.isolatedNameScopes[this.isolatedNameScopes.length - 1].recordDefinition(def);
}
parseOptionalSSAUseList(results) {
if (this.getToken().isNot(_.Token.percent_identifier)) {
return;
}
this.parseCommaSeparatedList('none', () => {
const result = this.parseSSAUse();
results.push(result);
});
}
parseSSAUse(allowResultNumber = true) {
const name = this.getTokenSpelling().str();
this.parseToken(_.Token.percent_identifier, "Expected SSA operand");
let number = 0;
if (this.getToken().is(_.Token.hash_identifier)) {
if (!allowResultNumber) {
this.emitError('Result number not allowed in argument list');
}
const value = this.getTokenSpelling().str().substring(1);
this.consumeToken(_.Token.hash_identifier);
number = parseInt(value, 10);
if (isNaN(number)) {
this.emitError(`Invalid SSA value result number '${value}'`);
}
}
return new _.UnresolvedOperand(null, name, number);
}
resolveSSAUse(unresolvedOperand, type) {
if (type !== null && type instanceof _.Type === false) {
throw new mlir.Error(`Type expected instead of '${type}'.`);
}
if (unresolvedOperand instanceof _.UnresolvedOperand) {
const entries = this.getSSAValueEntry(unresolvedOperand.name);
if (unresolvedOperand.number < entries.length && entries[unresolvedOperand.number]) {
const entry = entries[unresolvedOperand.number];
if (type && entry.value) {
entry.value.type = type;
}
return entry.value;
}
const value = new _.Value(unresolvedOperand.toString(), type);
if (entries.length <= unresolvedOperand.number) {
entries.length = unresolvedOperand.number + 1;
}
entries[unresolvedOperand.number] = { value, loc: unresolvedOperand.location };
return value;
}
// Handle literal operands (e.g., integer literals for buildable types)
if (unresolvedOperand && unresolvedOperand.literal) {
return new _.Value(unresolvedOperand.name, type);
}
this.emitError(`UnresolvedOperand expected, got '${JSON.stringify(unresolvedOperand)}'`);
return null;
}
getSSAValueEntry(name) {
const scope = this.isolatedNameScopes[this.isolatedNameScopes.length - 1];
if (!scope.values.has(name)) {
scope.values.set(name, []);
}
return scope.values.get(name);
}
parseBlock(block) {
block.operations = Array.isArray(block.operations) ? block.operations : [];
block.arguments = Array.isArray(block.arguments) ? block.arguments : [];
this.parseToken(_.Token.l_brace, "expected '{' to begin a region");
if (this.getToken().kind === _.Token.caret_identifier || (this.getToken().kind === _.Token.bare_identifier && this.getToken().getSpelling().str() && this.getToken().getSpelling().str().startsWith('^'))) {
if (this.getToken().kind === _.Token.caret_identifier) {
block.name = this.getTokenSpelling().str();
this.consumeToken(_.Token.caret_identifier);
} else {
block.name = this.getTokenSpelling().str();
this.consumeToken(_.Token.bare_identifier);
}
if (this.consumeIf(_.Token.l_paren)) {
while (!this.consumeIf(_.Token.r_paren) && this.getToken().isNot(_.Token.caret_identifier)) {
const value = this.getTokenSpelling().str();
this.consumeToken(_.Token.percent_identifier);
this.consumeToken(_.Token.colon);
const type = this.parseType();
const arg = { value, type };
const loc = this.parseTrailingLocationSpecifier();
if (loc) {
arg.loc = loc;
}
block.arguments.push(arg);
this.consumeIf(_.Token.comma);
}
}
if (block.name && block.name.endsWith(':')) {
block.name = block.name.slice(0, -1);
} else {
this.consumeToken(_.Token.colon);
}
}
while (!this.consumeIf(_.Token.r_brace)) {
if (this.getToken().kind === _.Token.caret_identifier || (this.getToken().kind === _.Token.bare_identifier && this.getToken().getSpelling().str() && this.getToken().getSpelling().str().startsWith('^'))) {
break;
}
const op = this.parseOperation();
block.operations.push(op);
}
block.loc = this.parseTrailingLocationSpecifier();
return block;
}
parseRegion(region, entryArguments, isIsolatedNameScope) {
// Push a new name scope for this region
this.pushSSANameScope(isIsolatedNameScope || false);
region.blocks = Array.isArray(region.blocks) ? region.blocks : [];
// Register SSA entries for entry arguments BEFORE parsing the block
// This ensures operations that reference %arg0 find the pre-registered entries
const resolvedEntryArgs = [];
if (entryArguments && entryArguments.length > 0) {
for (let i = 0; i < entryArguments.length; i++) {
const arg = entryArguments[i];
// Use explicit name if provided, otherwise generate %arg0, %arg1, etc.
const name = arg.name || `%arg${i}`;
const operand = new _.Value(name, arg.type);
// Register in SSA scope so operations can find it
this.addDefinition({ name, number: 0, location: arg.loc || null }, operand);
resolvedEntryArgs.push(operand);
}
}
const block = {};
this.parseBlock(block);
if (resolvedEntryArgs.length > 0) {
if (block.arguments.length === 0) {
block.arguments = resolvedEntryArgs;
} else if (block.arguments.length !== resolvedEntryArgs.length) {
this.emitError(`Entry block has ${block.arguments.length} arguments, but function signature has ${resolvedEntryArgs.length}`);
}
}
region.blocks.push(block);
let hasMultipleBlocks = false;
while ((this.getToken().kind === _.Token.caret_identifier || (this.getToken().kind === _.Token.bare_identifier && this.getToken().getSpelling().str() && this.getToken().getSpelling().str().startsWith('^'))) && this.getToken().isNot(_.Token.r_brace)) {
hasMultipleBlocks = true;
const nextBlock = {};
nextBlock.operations = [];
nextBlock.arguments = [];
if (this.getToken().kind === _.Token.caret_identifier) {
nextBlock.name = this.getTokenSpelling().str();
this.consumeToken(_.Token.caret_identifier);
} else {
nextBlock.name = this.getTokenSpelling().str();
this.consumeToken(_.Token.bare_identifier);
}
if (this.consumeIf(_.Token.l_paren)) {
while (!this.consumeIf(_.Token.r_paren)) {
const value = this.getTokenSpelling().str();
this.consumeToken(_.Token.percent_identifier);
this.consumeToken(_.Token.colon);
const type = this.parseType();
const arg = { value, type };
const loc = this.parseTrailingLocationSpecifier();
if (loc) {
arg.loc = loc;
}
nextBlock.arguments.push(arg);
this.consumeIf(_.Token.comma);
}
}
if (nextBlock.name && nextBlock.name.endsWith(':')) {
nextBlock.name = nextBlock.name.slice(0, -1);
} else {
this.consumeToken(_.Token.colon);
}
while (!(this.getToken().kind === _.Token.caret_identifier || (this.getToken().kind === _.Token.bare_identifier && this.getTokenSpelling().str() && this.getTokenSpelling().str().startsWith('^'))) && this.getToken().isNot(_.Token.r_brace)) {
const op = this.parseOperation();
nextBlock.operations.push(op);
}
region.blocks.push(nextBlock);
}
if (hasMultipleBlocks) {
this.consumeIf(_.Token.r_brace);
}
this.popSSANameScope();
return region;
}
parseTrailingLocationSpecifier() {
if (!this.consumeIf(_.Token.kw_loc)) {
return null;
}
this.parseToken(_.Token.l_paren, "expected '(' in location");
const tok = this.getToken();
let directLoc = null;
if (tok.is(_.Token.hash_identifier) && !tok.getSpelling().str().includes('.')) {
directLoc = this.parseLocationAlias();
} else {
directLoc = this.parseLocationInstance();
}
this.parseToken(_.Token.r_paren, "expected ')' in location");
return directLoc;
}
redirect(name) {
if (name === 'func.constant' && this.getToken().isNot(_.Token.at_identifier)) {
// Redirect for old std.constant should be arith.constant
name = 'arith.constant';
} else if (name.startsWith('check.')) {
// Redirect for stablehlo and iree both registering incompatible `check.` ops. Peek the next token to disambiguate.
const dialect = this.getToken().is(_.Token.l_paren) || this.getToken().is(_.Token.less) ? 'iree' : 'stablehlo';
name = name.replace('check.', `check.<${dialect}>.`);
} else if (this._redirect.has(name)) {
name = this._redirect.get(name);
}
return name;
}
};
_.AsmParser = class {
getNameLoc() {
return this.nameLoc;
}
parseKeyword(keyword) {
if (!this.parser.isCurrentTokenAKeyword()) {
this.emitError(`Expected '${keyword || 'keyword'}'`);
}
if (keyword && this.parser.getTokenSpelling().str() !== keyword) {
this.emitError(`Expected '${keyword}'`);
}
const spelling = this.parser.getTokenSpelling().str();
this.parser.consumeToken();
return spelling;
}
parseOptionalKeyword(arg) {
if (typeof arg === 'string') {
const keyword = arg;
if (!this.parser.isCurrentTokenAKeyword() || this.parser.getTokenSpelling().str() !== keyword) {
return false;
}
this.parser.consumeToken();
return true;
}
if (Array.isArray(arg)) {
const allowedValues = arg;
if (this.parser.getToken().is(_.Token.bare_identifier) || this.parser.getToken().isKeyword() || this.parser.getToken().is(_.Token.inttype)) {
const value = this.parser.getTokenSpelling().str();
if (allowedValues === undefined || allowedValues.some((v) => v === value)) {
this.parser.consumeToken();
return value;
}
}
return null;
}
if (arg === undefined) {
return this.parser.parseOptionalKeyword();
}
this.emitError('Invalid optional keyword');
return null;
}
parseKeywordType(keyword) {
this.parseKeyword(keyword);
return this.parser.parseType();
}
parseTypeList() {
return this.parser.parseTypeListNoParens();
}
parseType() {
return this.parser.parseType();
}
parseOptionalType() {
return this.parser.parseOptionalType();
}
parseDimensionList(allowDynamic, withTrailingX) {
return this.parser.parseDimensionListRanked(allowDynamic, withTrailingX);
}
parseXInDimensionList() {
return this.parser.parseXInDimensionList();
}
parseColonType() {
this.parser.parseToken(_.Token.colon, "expected ':'");
return this.parser.parseType();
}
parseColonTypeList() {
this.parser.parseToken(_.Token.colon, "expected ':'");
return this.parseTypeList();
}
parseOptionalColonTypeList() {
if (this.parser.consumeIf(_.Token.colon)) {
return this.parseTypeList();
}
return [];
}
parseArrowTypeList() {
this.parseArrow();
return this.parser.parseFunctionResultTypes();
}
parseOptionalArrowTypeList() {
if (this.parser.consumeIf(_.Token.arrow)) {
return this.parser.parseFunctionResultTypes();
}
return [];
}
parseAttribute(type) {
return this.parser.parseAttribute(type);
}
parseOptionalAttribute(type) {
return this.parser.parseOptionalAttribute(type);
}
parseCustomAttributeWithFallback(attrT, type) {
if (attrT) {
return attrT(this, type);
}
return this.parser.parseAttribute(type);
}
parseCustomTypeWithFallback(typeT) {
if (typeT && this.parser.getToken().isNot(_.Token.exclamation_identifier)) {
if (typeof typeT === 'function') {
return typeT(this);
}
const index = typeT.name.indexOf('.');
if (index === -1) {
throw new mlir.Error(`Invalid type name '${typeT.name}.`);
}
const dialectName = typeT.name.substring(0, index);
const dialect = this.parser.context.getOrLoadDialect(dialectName);
if (!dialect) {
if (this.parser.context.undocumented(dialectName)) {
this.emitError(`Undocumented dialect '${dialectName}'`);
} else {
this.emitError(`Unsupported dialect '${dialectName}'`);
}
}
return dialect.parseCustomTypeWithFallback(this, typeT.type);
}
return this.parser.parseType();
}
parseOptionalAttrDict(attributes) {
if (this.parser.getToken().is(_.Token.l_brace)) {
this.parser.parseAttributeDict(attributes);
}
}
parseOptionalAttrDictWithKeyword(attributes) {
if (this.parseOptionalKeyword('attributes')) {
this.parser.parseAttributeDict(attributes);
}
}
parseOptionalString() {
return this.parser.parseOptionalString();
}
parseOptionalKeywordOrString() {
return this.parser.parseOptionalKeywordOrString();
}
parseString() {
const value = this.parser.parseOptionalString();
if (value === null) {
this.emitError('Expected string');
}
return value;
}
parseOptionalInteger(type) {
return this.parser.parseOptionalInteger(type);
}
parseOptionalDecimalInteger() {
return this.parser.parseOptionalDecimalInteger();
}
parseInteger(type) {
return this.parser.parseInteger(type);
}
parseDecimalInteger() {
return this.parser.parseDecimalInteger();
}
parseCommaSeparatedList(delimiter, parseElement) {
return this.parser.parseCommaSeparatedList(delimiter, parseElement);
}
parseArrow() {
this.parser.parseToken(_.Token.arrow, "expected '->'");
}
parseOptionalArrow() {
return this.parser.consumeIf(_.Token.arrow);
}
parseEqual() {
this.parser.parseToken(_.Token.equal, "expected '='");
}
parseOptionalEqual() {
return this.parser.consumeIf(_.Token.equal);
}
parseLBrace() {
this.parser.parseToken(_.Token.l_brace, "expected '{'");
}
parseRBrace() {
this.parser.parseToken(_.Token.r_brace, "expected '}'");
}
parseOptionalLBrace() {
return this.parser.consumeIf(_.Token.l_brace);
}
parseOptionalRBrace() {
return this.parser.consumeIf(_.Token.r_brace);
}
parseComma() {
this.parser.parseToken(_.Token.comma, "expected ','");
}
parseOptionalComma() {
return this.parser.consumeIf(_.Token.comma);
}
parseLParen() {
this.parser.parseToken(_.Token.l_paren, "expected '('");
}
parseRParen() {
this.parser.parseToken(_.Token.r_paren, "expected ')'");
}
parseOptionalLParen() {
return this.parser.consumeIf(_.Token.l_paren);
}
parseOptionalRParen() {
return this.parser.consumeIf(_.Token.r_paren);
}
parseLSquare() {
return this.parser.parseToken(_.Token.l_square, "expected '['");
}
parseRSquare() {
return this.parser.parseToken(_.Token.r_square, "expected ']'");
}
parseOptionalLSquare() {
return this.parser.consumeIf(_.Token.l_square);
}
parseOptionalRSquare() {
return this.parser.consumeIf(_.Token.r_square);
}
parseColon() {
this.parser.parseToken(_.Token.colon, "expected ':'");
}
parseOptionalColon() {
return this.parser.consumeIf(_.Token.colon);
}
parseLess() {
this.parser.parseToken(_.Token.less, "expected '<'");
}
parseGreater() {
this.parser.parseToken(_.Token.greater, "expected '>'");
}
parseOptionalLess() {
return this.parser.consumeIf(_.Token.less);
}
parseOptionalGreater() {
return this.parser.consumeIf(_.Token.greater);
}
parseOptionalVerticalBar() {
return this.parser.consumeIf(_.Token.vertical_bar);
}
parseOptionalQuestion() {
return this.parser.consumeIf(_.Token.question);
}
parseOptionalStar() {
return this.parser.consumeIf(_.Token.star);
}
parseOptionalMinus() {
return this.parser.consumeIf(_.Token.minus);
}
parseOptionalEllipsis() {
return this.parser.consumeIf(_.Token.ellipsis);
}
parseOptionalPlus() {
return this.parser.consumeIf(_.Token.plus);
}
parseStar() {
this.parser.parseToken(_.Token.star, "expected '*'");
}
parsePlus() {
this.parser.parseToken(_.Token.plus, "expected '+'");
}
parseQuestion() {
this.parser.parseToken(_.Token.question, "expected '?'");
}
parseVerticalBar() {
this.parser.parseToken(_.Token.vertical_bar, "expected '|'");
}
parseEllipsis() {
this.parser.parseToken(_.Token.ellipsis, "expected '...'");
}
parseMinus() {
this.parser.parseToken(_.Token.minus, "expected '-'");
}
emitError(...args) {
return this.parser.emitError(...args);
}
getCurrentLocation() {
return this.parser.getToken().loc;
}
parseSymbolName(attrName, attrs) {
const result = this.parseOptionalSymbolName();
if (result === null) {
this.emitError(`Expected valid '@'-identifier for symbol name`);
}
attrs.set(attrName, result);
}
parseOptionalSymbolName() {
const atToken = this.parser.getToken();
if (atToken.isNot(_.Token.at_identifier)) {
return null;
}
const result = new _.StringAttr(atToken.getSymbolReference());
this.parser.consumeToken();
if (this.parser.state.asmState) {
// parser.state.asmState.addUses(SymbolRefAttr::get(result), atToken.getLocRange())
}
return result;
}
parseAffineMapOfSSAIds(mapOperands, mapAttr, attrName, attrs, delimiter = 'Square') {
const dimOperands = [];
const symOperands = [];
const parseElement = (isSymbol) => {
const operand = this.parseOperand();
if (isSymbol) {
symOperands.push(operand);
} else {
dimOperands.push(operand);
}
};
const map = this.parser.parseAffineMapOfSSAIds(parseElement, delimiter);
if (map) {
attrs.set(attrName, map);
}
mapOperands.push(...dimOperands, ...symOperands);
}
parseAffineMap() {
return this.parser.parseAffineMapReference();
}
parseAffineExprOfSSAIds(mapOperands) {
const dimOperands = [];
const symOperands = [];
const parseElement = (isSymbol) => {
const operand = this.parseOperand();
if (isSymbol) {
symOperands.push(operand);
} else {
dimOperands.push(operand);
}
};
const expr = this.parser.parseAffineExprOfSSAIds(parseElement);
mapOperands.push(...dimOperands, ...symOperands);
return expr;
}
};
_.OpAsmParser = class extends _.AsmParser {
parseFunctionOp(op, allowVariadic) {
this.parseOptionalVisibilityKeyword(op.attributes);
this.parseSymbolName('sym_name', op.attributes);
const sig = this.parseFunctionSignatureWithArguments(allowVariadic);
const argTypes = [];
for (const arg of sig.arguments) {
if (arg.name !== '...') {
argTypes.push(arg.type);
}
}
const type = new _.FunctionType(argTypes, sig.resultTypes);
op.addAttribute('function_type', new _.TypeAttrOf(type));
if (sig.resultAttrs.some((a) => a !== null)) {
op.addAttribute('res_attrs', sig.resultAttrs);
}
const argAttrs = sig.arguments.filter((a) => a.name !== '...').map((a) => a.attrs || null);
if (argAttrs.some((a) => a !== null)) {
op.addAttribute('arg_attrs', argAttrs);
}
this.parseOptionalAttrDictWithKeyword(op.attributes);
if (this.parser.getToken().is(_.Token.l_brace)) {
const region = op.addRegion();
// Functions are IsolatedFromAbove, so pass true for isIsolatedNameScope
this.parseRegion(region, sig.arguments, /* isIsolatedNameScope */ true);
}
}
parseFunctionSignature(argOperands) {
const argTypes = [];
const argAttrs = [];
const resultTypes = [];
const resultAttrs = [];
this.parser.parseToken(_.Token.l_paren, "expected '(' in function signature");
if (this.parser.getToken().isNot(_.Token.r_paren)) {
this.parseTypeAndAttrList(argTypes, argAttrs, argOperands);
}
this.parser.parseToken(_.Token.r_paren, "expected ')' in function signature");
if (this.parser.consumeIf(_.Token.arrow)) {
this.parseFunctionResultList(resultTypes, resultAttrs);
}
return { argTypes, argAttrs, resultTypes, resultAttrs };
}
parseFunctionSignatureWithArguments(allowVariadic) {
const argResult = this.parseFunctionArgumentList(allowVariadic);
const resultTypes = [];
const resultAttrs = [];
if (this.parser.consumeIf(_.Token.arrow)) {
this.parseFunctionResultList(resultTypes, resultAttrs);
}
return { arguments: argResult.arguments, isVariadic: argResult.isVariadic, resultTypes, resultAttrs };
}
parseFunctionResultList(types, attrs) {
if (this.parser.consumeIf(_.Token.l_paren)) {
if (this.parser.consumeIf(_.Token.r_paren)) {
return;
}
this.parseTypeAndAttrList(types, attrs);
this.parser.parseToken(_.Token.r_paren, "expected ')' in function result list");
} else {
const type = this.parseType();
types.push(type);
attrs.push(null);
}
}
// Returns { arguments: Array<OpAsmParser.Argument>, isVariadic: boolean }
parseFunctionArgumentList(allowVariadic) {
const inputs = [];
let isVariadic = false;
if (this.parser.consumeIf(_.Token.l_paren)) {
while (!this.parser.consumeIf(_.Token.r_paren)) {
if (this.parser.getToken().is(_.Token.r_paren)) {
break;
}
if (allowVariadic && this.parser.consumeIf(_.Token.ellipsis)) {
isVariadic = true;
this.parser.parseToken(_.Token.r_paren, "expected ')' after '...'");
break;
}
if (this.parser.getToken().is(_.Token.percent_identifier)) {
const ssaName = this.parseOperand();
this.parser.consumeToken(_.Token.colon);
const type = this.parseType();
let attrs = null;
if (this.parser.getToken().is(_.Token.l_brace)) {
attrs = new Map();
this.parser.parseAttributeDict(attrs);
}
const loc = this.parseOptionalLocationSpecifier();
inputs.push(new _.OpAsmParser.Argument(ssaName, type, attrs, loc));
} else {
// Type-only argument (no explicit name like %arg0)
// Don't generate a name - let the region/SSA system handle it
const type = this.parseType();
let attrs = null;
if (this.parser.getToken().is(_.Token.l_brace)) {
attrs = new Map();
this.parser.parseAttributeDict(attrs);
}
inputs.push(new _.OpAsmParser.Argument(null, type, attrs, null));
}
if (this.parser.getToken().isNot(_.Token.r_paren)) {
if (!this.parser.consumeIf(_.Token.comma)) {
break;
}
if (this.parser.getToken().is(_.Token.r_paren)) {
break;
}
}
}
}
return { arguments: inputs, isVariadic };
}
parseTypeAndAttrList(types, attrs, operands) {
let index = 0;
this.parseCommaSeparatedList('none', () => {
const type = this.parseType();
types.push(type);
// Parse optional attribute dict after each type
if (this.parser.getToken().is(_.Token.l_brace)) {
const attrList = new Map();
this.parser.parseAttributeDict(attrList);
attrs.push(attrList);
// Associate attrs with operand if available
if (operands && index < operands.length) {
operands[index].attributes = attrList;
}
} else {
attrs.push(null);
}
index++;
return true;
});
}
parseDenseI64ArrayAttr(attributeName, attributes) {
this.parseKeyword(attributeName);
this.parseEqual();
this.parseLSquare();
const values = [];
if (!this.parseOptionalRSquare()) {
do {
values.push(this.parseInteger());
} while (this.parseOptionalComma());
this.parseRSquare();
}
attributes.set(attributeName, new _.DenseI64ArrayAttr(values));
}
parseOptionalVisibilityKeyword(attributes) {
const visibility = this.parseOptionalKeyword(['public', 'private', 'nested']);
if (visibility) {
attributes.set('sym_visibility', visibility);
}
}
parseOperands(operands) {
this.parseCommaSeparatedList('none', () => {
operands.push(this.parseOperand());
});
}
parseShapedResultList(operands, operandTypes, resultTypes, tiedOperands, resolvedOperands) {
const indexType = new _.IndexType();
const tiedOperandIndices = [];
do {
let type = null;
let tiedOperandIndex = -1;
const tiedResult = this.parseOptionalOperand();
if (tiedResult) {
tiedOperandIndex = this.findTiedOperand(tiedResult, operands);
if (this.parseOptionalKeyword('as')) {
type = this.parseType();
} else if (tiedOperandIndex >= 0 && tiedOperandIndex < operandTypes.length) {
type = operandTypes[tiedOperandIndex];
}
} else {
type = this.parseType();
}
const dims = this.parseOperandList('optionalBrace');
if (resolvedOperands) {
this.resolveOperands(dims, dims.map(() => indexType), resolvedOperands);
}
if (type) {
resultTypes.push(type);
}
tiedOperandIndices.push(tiedOperandIndex);
} while (this.parseOptionalComma());
if (tiedOperands && tiedOperandIndices.length > 0) {
tiedOperands.push(...tiedOperandIndices);
}
}
findTiedOperand(tiedResult, operands) {
for (let i = 0; i < operands.length; i++) {
if (operands[i].name === tiedResult.name && operands[i].number === tiedResult.number) {
return i;
}
}
return -1;
}
};
_.OpAsmParser.Argument = class {
constructor(ssaName, type, attrs, loc) {
this.ssaName = ssaName;
this.type = type;
this.attrs = attrs;
this.loc = loc;
}
get name() {
return this.ssaName ? this.ssaName.name : null;
}
get value() {
return this.ssaName ? this.ssaName.toString() : null;
}
};
_.CustomOpAsmParser = class extends _.OpAsmParser {
constructor(resultIDs, parser) {
super();
this.resultIDs = resultIDs || [];
this.parser = parser;
this.nameLoc = parser.getToken().loc.copy();
}
parseOperation() {
return this.parser.parseOperation();
}
getNumResults() {
let count = 0;
for (const entry of this.resultIDs) {
count += entry[1];
}
return count;
}
getResultName(index) {
if (index < this.resultIDs.length) {
return this.resultIDs[index][0];
}
return null;
}
parseArgument(allowType, allowAttrs) {
const ssaName = this.parseOperand();
let type = null;
let attrs = null;
let loc = null;
if (allowType) {
type = this.parseColonType();
}
if (allowAttrs) {
attrs = {};
this.parseOptionalAttrDict(attrs);
if (Object.keys(attrs).length === 0) {
attrs = null;
}
}
loc = this.parseOptionalLocationSpecifier();
return new _.OpAsmParser.Argument(ssaName, type, attrs, loc);
}
parseOptionalArgument(allowType = false, allowAttrs = false) {
if (this.parser.getToken().is(_.Token.percent_identifier)) {
return this.parseArgument(allowType, allowAttrs);
}
return null;
}
parseArgumentList(delimiter, allowType = false, allowAttrs = false) {
delimiter = delimiter || 'none';
if (delimiter === 'none') {
if (this.parser.getToken().isNot(_.Token.percent_identifier)) {
return [];
}
}
const parseOneArgument = () => {
if (this.parser.getToken().is(_.Token.percent_identifier)) {
return this.parseArgument(allowType, allowAttrs);
}
return null;
};
return this.parseCommaSeparatedList(delimiter, parseOneArgument);
}
parseAssignmentList(lhs, rhs) {
const result = this.parseOptionalAssignmentList(lhs, rhs);
if (!result) {
this.emitError(`Expected '('`);
}
}
parseOptionalAssignmentList(lhs, rhs) {
if (!this.parseOptionalLParen()) {
return false;
}
if (!this.parseOptionalRParen()) {
do {
lhs.push(this.parseArgument());
this.parseEqual();
rhs.push(this.parseOperand());
} while (this.parseOptionalComma());
this.parseRParen();
}
return true;
}
parseRegion(region, entryArguments, enableNameShadowing) {
return this.parser.parseRegion(region, entryArguments, enableNameShadowing);
}
parseOptionalRegion(region, entryArguments, enableNameShadowing) {
if (this.parser.getToken().isNot(_.Token.l_brace)) {
region.blocks = region.blocks || [];
return false;
}
return this.parser.parseRegion(region, entryArguments, enableNameShadowing);
}
parseSuccessor() {
if (this.parser.getToken().isNot(_.Token.caret_identifier)) {
this.emitError('Expected block name');
}
const label = this.parser.getTokenSpelling().str();
this.parser.consumeToken(_.Token.caret_identifier);
return { label };
}
parseOptionalSuccessor() {
if (this.parser.getToken().isNot(_.Token.caret_identifier)) {
return null;
}
return this.parseSuccessor();
}
parseSuccessorAndUseList() {
const successor = this.parseSuccessor();
if (this.parseOptionalLParen()) {
const operands = this.parseOperandList('none');
this.parseRParen();
successor.operands = operands;
}
return successor;
}
parseOperand(allowResultNumber = true) {
return this.parser.parseSSAUse(allowResultNumber);
}
parseOptionalOperand(allowResultNumber = true) {
if (this.parser.getToken().is(_.Token.percent_identifier)) {
return this.parseOperand(allowResultNumber);
}
return null;
}
parseOperandList(delimiter) {
delimiter = delimiter || 'none';
if (delimiter === 'none') {
if (this.parser.getToken().isNot(_.Token.percent_identifier)) {
return [];
}
}
const parseOneOperand = () => {
if (this.parser.getToken().is(_.Token.percent_identifier)) {
return this.parseOperand();
}
return null;
};
return this.parseCommaSeparatedList(delimiter, parseOneOperand);
}
parseTrailingOperandList(delimiter) {
if (!this.parseOptionalComma()) {
return [];
}
return this.parseOperandList(delimiter);
}
parseDimAndSymbolList(operands) {
const indexType = new _.IndexType();
const opInfos = this.parseOperandList('paren');
const numDims = opInfos.length;
const symInfos = this.parseOperandList('optionalSquare');
opInfos.push(...symInfos);
this.resolveOperands(opInfos, opInfos.map(() => indexType), operands);
return numDims;
}
resolveOperand(operand, type, result) {
const resolved = this.parser.resolveSSAUse(operand, type);
if (result) {
result.push(resolved);
}
return resolved;
}
resolveOperands(operands, types, result) {
if (!Array.isArray(operands)) {
throw new mlir.Error(`Unexpected operands type '${typeof operands}'.`);
}
if (!Array.isArray(types)) {
return;
}
const count = Math.min(operands.length, types.length);
if (result) {
for (let i = 0; i < count; i++) {
const operand = operands[i];
const type = types[i];
const resolved = this.parser.resolveSSAUse(operand, type);
result.push(resolved);
}
} else {
for (let i = 0; i < count; i++) {
const operand = operands[i];
const type = types[i];
if (operand && type) {
if (type instanceof _.Type === false) {
throw new mlir.Error(`Type expected instead of '${type}'.`);
}
if (operand instanceof _.Value) {
operand.type = type;
} else if (typeof operand === 'object') {
operand.type = type;
}
}
}
}
}
parseGenericOperation() {
return this.parser.parseGenericOperation();
}
parseCustomOperationName() {
return this.parser.parseCustomOperationName();
}
parseOptionalLocationSpecifier() {
// Ref impl: CustomOpAsmParser::parseOptionalLocationSpecifier (Parser.cpp:2002)
// Separate implementation from parseTrailingLocationSpecifier
// If there is a 'loc' we parse a trailing location.
if (!this.parser.consumeIf(_.Token.kw_loc)) {
return null;
}
this.parser.parseToken(_.Token.l_paren, "expected '(' in location");
const tok = this.parser.getToken();
let directLoc = null;
// Check to see if we are parsing a location alias. We are parsing a
// location alias if the token is a hash identifier *without* a dot in it -
// the dot signifies a dialect attribute. Otherwise, we parse the location
// directly.
if (tok.is(_.Token.hash_identifier) && !tok.getSpelling().str().includes('.')) {
directLoc = this.parser.parseLocationAlias();
} else {
directLoc = this.parser.parseLocationInstance();
}
this.parser.parseToken(_.Token.r_paren, "expected ')' in location");
return directLoc;
}
};
_.DialectAsmParser = class extends _.AsmParser {
};
_.CustomDialectAsmParser = class extends _.DialectAsmParser {
constructor(fullSpec, parser) {
super();
this.fullSpec = fullSpec;
this.parser = parser;
this.nameLoc = parser.getToken().loc.copy();
}
getFullSymbolSpec() {
return this.fullSpec.str();
}
};
_.TensorLiteralParser = class {
constructor(parser) {
this._parser = parser;
// Reference stores pairs of (isNegative, token) for proper type validation
// We store {isNegative, value, kind} objects
this._storage = [];
this._shape = [];
this._hexStorage = null;
}
parse(allowHex) {
if (allowHex && this._parser.getToken().is(_.Token.string)) {
const hexStr = this._parser.getToken().getStringValue();
this._parser.consumeToken(_.Token.string);
if (hexStr.startsWith('0x')) {
const cleanHex = hexStr.substring(2);
const data = new Uint8Array(cleanHex.length >> 1);
for (let i = 0; i < data.length; i++) {
const index = i << 1;
data[i] = parseInt(cleanHex.substring(index, index + 2), 16);
}
this._hexStorage = data;
return { storage: data, shape: null };
}
// Non-hex string element
this._storage.push({ isNegative: false, value: hexStr, kind: 'string' });
return { storage: this._storage, shape: this._shape };
}
if (this._parser.getToken().is(_.Token.l_square)) {
this.parseList(this._shape);
} else {
this.parseElement();
// Single element parsed without list - shape stays empty (splat)
}
return { storage: this._storage, shape: this._shape };
}
parseList(dims) {
this._parser.consumeToken(_.Token.l_square);
let first = true;
let newDims = [];
let size = 0;
while (!this._parser.consumeIf(_.Token.r_square)) {
const thisDims = [];
if (this._parser.getToken().is(_.Token.l_square)) {
this.parseList(thisDims);
} else {
this.parseElement();
}
size++;
if (!first) {
const compareDims = (a, b) =>{
if (a.length !== b.length) {
return false;
}
for (let i = 0; i < a.length; i++) {
if (a[i] !== b[i]) {
return false;
}
}
return true;
};
// Verify consistent dimensions (reference checks prevDims == newDims)
const dimsMatch = compareDims(thisDims, newDims);
if (!dimsMatch) {
this._parser.emitError('Tensor literal is invalid; ranks are not consistent between elements');
}
}
newDims = thisDims;
first = false;
this._parser.consumeIf(_.Token.comma);
}
dims.length = 0;
dims.push(size);
dims.push(...newDims);
}
parseElement() {
switch (this._parser.getToken().kind) {
// Parse a boolean element.
case _.Token.kw_true:
case _.Token.kw_false: {
const value = this._parser.getTokenSpelling().str();
this._parser.consumeToken(this._parser.getToken().kind);
this._storage.push({ isNegative: false, value, kind: 'boolean' });
break;
}
case _.Token.floatliteral:
case _.Token.integer: {
const value = this._parser.getTokenSpelling().str();
const kind = this._parser.getToken().kind === _.Token.floatliteral ? 'float' : 'int';
this._parser.consumeToken(this._parser.getToken().kind);
this._storage.push({ isNegative: false, value, kind });
break;
}
case _.Token.minus: {
this._parser.consumeToken(_.Token.minus);
if (!this._parser.getToken().isAny(_.Token.floatliteral, _.Token.integer)) {
this._parser.emitError(`Expected integer or floating point literal`);
}
const value = this._parser.getTokenSpelling().str();
const kind = this._parser.getToken().kind === _.Token.floatliteral ? 'float' : 'int';
this._parser.consumeToken(this._parser.getToken().kind);
this._storage.push({ isNegative: true, value, kind });
break;
}
case _.Token.string: {
const value = this._parser.getToken().getStringValue();
this._parser.consumeToken(_.Token.string);
this._storage.push({ isNegative: false, value, kind: 'string' });
break;
}
case _.Token.l_paren:
this._parser.consumeToken(_.Token.l_paren);
this.parseElement();
this._parser.parseToken(_.Token.comma, "expected ',' between complex elements");
this.parseElement();
this._parser.parseToken(_.Token.r_paren, "expected ')' after complex elements");
break;
default:
this._parser.emitError('Expected element literal of primitive type');
}
}
getShape() {
return this._shape;
}
getAttr(type) {
if (this._hexStorage instanceof Uint8Array) {
return this._hexStorage;
}
const elementType = type && type.elementType ? type.elementType : null;
const numElements = type && type.getNumElements ? type.getNumElements().toNumber() : 0;
const isComplex = elementType instanceof _.ComplexType;
const baseElemType = isComplex && elementType.elementType ? elementType.elementType : elementType;
let isUnsigned = false;
let isInteger = false;
let isI1 = false;
let bitWidth = 0;
if (baseElemType) {
const typeStr = baseElemType.toString();
isUnsigned = typeStr.startsWith('ui');
isI1 = typeStr === 'i1';
const intMatch = typeStr.match(/^[su]?i(\d+)$/);
if (intMatch) {
isInteger = true;
bitWidth = parseInt(intMatch[1], 10);
} else if (typeStr === 'index') {
isInteger = true;
bitWidth = 64;
}
}
const convertElement = (elem) => {
const { isNegative, value, kind } = elem;
if (isNegative && isUnsigned) {
throw new mlir.Error(`Expected unsigned integer elements, but parsed negative value`);
}
if (kind === 'boolean') {
if (isInteger && !isI1) {
throw new mlir.Error(`Expected i1 type for 'true' or 'false' values`);
}
return value === true || value === 'true' ? 1 : 0;
}
if (kind === 'float' && isInteger) {
throw new mlir.Error(`Expected integer elements, but parsed floating-point`);
}
let result = null;
if (kind === 'int') {
if (!isInteger && baseElemType) {
// Integer literal for a float type → hex bit pattern (e.g., 0xFF800000 for -Infinity)
const bits = typeof value === 'string' && value.startsWith('0x') ? BigInt(value) : BigInt(isNegative ? `-${value}` : value);
const typeStr = baseElemType.toString();
if (typeStr === 'f32') {
const buf = new ArrayBuffer(4);
new DataView(buf).setUint32(0, Number(bits & 0xFFFFFFFFn));
result = new DataView(buf).getFloat32(0);
} else if (typeStr === 'f64') {
const buf = new ArrayBuffer(8);
new DataView(buf).setBigUint64(0, bits & 0xFFFFFFFFFFFFFFFFn);
result = new DataView(buf).getFloat64(0);
} else if (typeStr === 'f16') {
const buf = new ArrayBuffer(2);
new DataView(buf).setUint16(0, Number(bits & 0xFFFFn));
result = new DataView(buf).getFloat16(0);
} else if (typeStr === 'bf16') {
const buf = new ArrayBuffer(4);
new DataView(buf).setUint32(0, Number(bits & 0xFFFFn) << 16);
result = new DataView(buf).getFloat32(0);
} else {
result = Number(bits);
}
} else if (bitWidth >= 64) {
result = BigInt(isNegative ? `-${value}` : value);
} else {
result = Number(value);
if (isNegative) {
result = -result;
}
}
} else if (kind === 'float') {
result = parseFloat(value);
if (isNegative) {
result = -result;
}
} else {
result = value;
}
return result;
};
if (numElements === 0) {
return [];
}
const maxSplatExpansion = 10000000;
const isSplat = this._shape.length === 0 && this._storage.length > 0;
if (isSplat && numElements > 1) {
if (numElements > maxSplatExpansion) {
return null;
}
if (isComplex && this._storage.length === 2) {
const result = [];
const real = convertElement(this._storage[0]);
const imag = convertElement(this._storage[1]);
for (let i = 0; i < numElements; i++) {
result.push(new base.Complex(real, imag));
}
return result;
}
const converted = convertElement(this._storage[0]);
return new Array(numElements).fill(converted);
}
if (isComplex && Array.isArray(this._storage)) {
const result = [];
for (let i = 0; i < this._storage.length; i += 2) {
result.push(new base.Complex(convertElement(this._storage[i]), convertElement(this._storage[i + 1])));
}
return result;
}
return this._storage.map(convertElement);
}
};
_.EncodingReader = class {
constructor(data) {
if (data instanceof Uint8Array) {
this._data = data;
this._reader = base.BinaryReader.open(data);
} else {
this._data = null;
this._reader = data;
}
}
get length() {
return this._reader.length;
}
get position() {
return this._reader.position;
}
empty() {
return this.position >= this.length;
}
size() {
return this._reader.length - this._reader.position;
}
seek(offset) {
this._reader.seek(offset);
}
skipBytes(length) {
this._reader.skip(length);
}
parseBytes(length) {
return this._reader.read(length);
}
parseByte() {
return this._reader.byte();
}
peek() {
const position = this._reader.position;
const value = this._reader.byte();
this._reader.seek(position);
return value;
}
parseVarInt() {
let result = this._reader.byte();
if (result & 1) {
return BigInt(result >> 1);
}
if (result === 0) {
return this._reader.uint64();
}
result = BigInt(result);
let mask = 1n;
let numBytes = 0n;
let shift = 8n;
while (result > 0n && (result & mask) === 0n) {
if (numBytes >= 7n) {
throw new mlir.Error('Invalid varint.');
}
result |= (BigInt(this._reader.byte()) << shift);
mask <<= 1n;
shift += 8n;
numBytes++;
}
if (numBytes === 0n || numBytes > 7n) {
throw new mlir.Error(`Invalid varint.`);
}
result >>= numBytes + 1n;
return result;
}
parseSignedVarInt() {
const n = this.parseVarInt();
return (n >> 1n) ^ -(n & 1n);
}
parseVarIntWithFlag() {
const result = this.parseVarInt();
return [result >> 1n, (result & 1n) === 1n];
}
parseNullTerminatedString() {
const reader = this._reader;
let result = '';
let value = -1;
const maxLength = reader.length - reader.position;
let bytesRead = 0;
for (; ;) {
if (bytesRead >= maxLength) {
throw new mlir.Error('Malformed null-terminated string, no null character found.');
}
value = reader.byte();
bytesRead++;
if (value === 0x00) {
break;
}
result += String.fromCharCode(value);
}
return result;
}
parseEntry(entries, entryStr) {
const entryIdx = this.parseVarInt().toNumber();
return this.resolveEntry(entries, entryIdx, entryStr);
}
resolveEntry(entries, entryIdx, entryStr) {
if (entryIdx >= entries.length) {
throw new mlir.Error(`Invalid '${entryStr}' index.`);
}
return entries[entryIdx];
}
parseSection(alignmentValidator) {
const sectionIDAndHasAlignment = this.parseByte();
const length = this.parseVarInt().toNumber();
const sectionID = sectionIDAndHasAlignment & 0x7F;
const hasAlignment = sectionIDAndHasAlignment & 0x80;
if (sectionID >= 9) { // kNumSections
throw new mlir.Error(`Invalid section ID: ${sectionID}.`);
}
if (hasAlignment) {
const alignment = this.parseVarInt().toNumber();
alignmentValidator(alignment);
this.alignTo(alignment);
}
const sectionData = this.parseBytes(length);
return [sectionID, sectionData];
}
parseBlobAndAlignment() {
const alignment = this.parseVarInt().toNumber();
const dataSize = this.parseVarInt().toNumber();
this.alignTo(alignment);
const data = this.parseBytes(dataSize);
return [data, alignment];
}
alignTo(alignment) {
if ((alignment & (alignment - 1)) !== 0) {
throw new mlir.Error('Expected alignment to be a power-of-two.');
}
while ((this.position & (alignment - 1)) !== 0) {
const padding = this.parseByte();
if (padding !== 0xCB) {
throw new mlir.Error(`Expected alignment byte (0xCB), but got: 0x${padding.toString(16)}.`);
}
}
}
};
_.BytecodeDialect = class {
load(reader, ctx) {
this.dialect = ctx.getOrLoadDialect(this.name);
if (!this.dialect) {
throw new mlir.Error(`Dialect '${this.name}' is unknown.`);
}
this.interface = this.dialect;
if (this.versionBuffer) {
const encReader = new _.EncodingReader(this.versionBuffer, reader.loc);
const versionReader = reader.withEncodingReader(encReader);
const loadedVersion = this.interface.readVersion(versionReader);
if (!loadedVersion) {
return false;
}
}
return true;
}
getLoadedDialect() {
return this.dialect;
}
};
_.DialectBytecodeReader = class {
};
_.DialectReader = class extends _.DialectBytecodeReader {
constructor(attrTypeReader, stringReader, resourceReader, dialectsMap, reader, bytecodeVersion, depth = 0) {
super();
this.attrTypeReader = attrTypeReader;
this.stringReader = stringReader;
this.resourceReader = resourceReader;
this.dialectsMap = dialectsMap;
this.reader = reader;
this.bytecodeVersion = bytecodeVersion;
this.depth = depth;
this._floatBuffer = new ArrayBuffer(8);
this._floatView = new DataView(this._floatBuffer);
this._floatBitWidths = { f16: 16, bf16: 16, f32: 32, f64: 64, f80: 80, f128: 128, tf32: 19, f4E2M1FN: 4, f6E2M3FN: 6, f6E3M2FN: 6, f8E5M2: 8, f8E4M3: 8, f8E4M3FN: 8, f8E5M2FNUZ: 8, f8E4M3FNUZ: 8, f8E4M3B11FNUZ: 8, f8E3M4: 8, f8E8M0FNU: 8 };
}
readType() {
const index = this.reader.parseVarInt().toNumber();
if (this.attrTypeReader.isResolving() && this.depth > this.maxAttrTypeDepth) {
const existing = this.attrTypeReader.getTypeOrSentinel(index);
if (!existing) {
throw new mlir.Error(`Exceeded maximum type depth at '${index}'.`);
}
return existing;
}
return this.attrTypeReader.readType(index, this.depth + 1);
}
readAttribute() {
const index = this.reader.parseVarInt().toNumber();
if (this.attrTypeReader.isResolving() && this.depth > this.maxAttrTypeDepth) {
const existing = this.attrTypeReader.getAttributeOrSentinel(index);
if (!existing) {
throw new mlir.Error(`Exceeded maximum attribute depth at '${index}'.`);
}
return existing;
}
return this.attrTypeReader.readAttribute(index, this.depth + 1);
}
readOptionalAttribute() {
return this.attrTypeReader.parseOptionalAttribute(this.reader);
}
readString() {
return this.stringReader.parseString(this.reader);
}
readVarInt() {
return this.reader.parseVarInt().toNumber();
}
readSignedVarInt() {
return this.reader.parseSignedVarInt();
}
readByte() {
return this.reader.parseByte();
}
readBlob() {
const size = this.reader.parseVarInt().toNumber();
return this.reader.parseBytes(size);
}
readResourceHandle() {
return this.resourceReader.parseResourceHandle(this.reader);
}
readSignedVarInts() {
const count = this.reader.parseVarInt().toNumber();
const result = new Array(count);
for (let i = 0; i < count; i++) {
result[i] = this.reader.parseSignedVarInt();
}
return result;
}
readAPIntWithKnownWidth(bitWidth) {
if (bitWidth <= 8) {
return BigInt(this.reader.parseByte());
}
if (bitWidth <= 64) {
return this.reader.parseSignedVarInt();
}
const numWords = this.reader.parseVarInt().toNumber();
let value = 0n;
for (let i = 0; i < numWords; i++) {
const word = this.reader.parseSignedVarInt();
value |= (word << BigInt(i * 64));
}
return value;
}
readAPFloatWithKnownSemantics(type) {
const bitWidth = this._floatBitWidths[type.name];
if (!bitWidth) {
throw new mlir.Error(`Unsupported float type '${type.name}'.`);
}
const bits = this.readAPIntWithKnownWidth(bitWidth);
if (bitWidth <= 32) {
this._floatView.setUint32(0, Number(bits), true);
} else {
this._floatView.setBigUint64(0, bits, true);
}
switch (type.name) {
case 'f16': return this._floatView.getFloat16(0, true);
case 'f32': return this._floatView.getFloat32(0, true);
case 'f64': return this._floatView.getFloat64(0, true);
case 'bf16': return this._floatView.getBfloat16(0, true);
default: throw new mlir.Error(`Unsupported float type '${type.name}'.`);
}
}
withEncodingReader(encodingReader) {
return new _.DialectReader(this.attrTypeReader, this.stringReader, this.resourceReader, this.dialectsMap, encodingReader, this.bytecodeVersion, this.depth);
}
};
_.AttrTypeReader = class {
constructor(stringReader, resourceReader, dialectsMap, bytecodeVersion, fileLoc, config) {
this.stringReader = stringReader;
this.resourceReader = resourceReader;
this.dialectsMap = dialectsMap;
this.fileLoc = fileLoc;
this.bytecodeVersion = bytecodeVersion;
this.parserConfig = config;
this.attributes = [];
this.types = [];
this.maxAttrTypeDepth = 5;
this._resolving = false;
}
initialize(dialects, sectionData, offsetSectionData) {
this._sectionData = sectionData;
const offsetReader = new _.EncodingReader(offsetSectionData);
const numAttributes = offsetReader.parseVarInt().toNumber();
const numTypes = offsetReader.parseVarInt().toNumber();
this.attributes = new Array(numAttributes);
this.types = new Array(numTypes);
let currentOffset = 0;
const parseEntries = (entries) => {
let currentIndex = 0;
const endIndex = entries.length;
while (currentIndex < endIndex) {
const dialectIndex = offsetReader.parseVarInt().toNumber();
const dialect = dialects[dialectIndex];
const numEntries = offsetReader.parseVarInt().toNumber();
for (let i = 0; i < numEntries; i++) {
const entry = {};
const entrySizeWithFlag = offsetReader.parseVarInt();
entry.hasCustomEncoding = (entrySizeWithFlag & 1n) === 1n;
const entrySize = (entrySizeWithFlag >> 1n).toNumber();
// Store offset like old code (don't bound reading to entry size)
entry.offset = currentOffset;
entry.size = entrySize;
entry.dialect = dialect;
entry.resolved = null;
currentOffset += entrySize;
entries[currentIndex++] = entry;
}
}
};
// Process attributes, then types
parseEntries(this.attributes);
parseEntries(this.types);
}
isResolving() {
return this._resolving;
}
getAttributeOrSentinel(index) {
if (index >= this.attributes.length) {
return null;
}
return this.attributes[index].resolved;
}
getTypeOrSentinel(index) {
if (index >= this.types.length) {
return null;
}
return this.types[index].resolved;
}
readAttribute(index, depth = 0) {
return this.readEntry(this.attributes, index, 'attr', depth);
}
readType(index, depth = 0) {
return this.readEntry(this.types, index, 'type', depth);
}
resolveAttribute(index, depth = 0) {
// Resolve an attribute at the given index (equivalent to resolveEntry in C++)
return this.resolveEntry(this.attributes, index, 'attr', depth);
}
resolveType(index, depth = 0) {
// Resolve a type at the given index (equivalent to resolveEntry in C++)
return this.resolveEntry(this.types, index, 'type', depth);
}
parseAttribute(reader) {
// Parse an attribute reference from the reader (varint index)
const index = reader.parseVarInt().toNumber();
return this.resolveAttribute(index);
}
parseOptionalAttribute(reader) {
const [attrIdx, flag] = reader.parseVarIntWithFlag();
if (!flag) {
return null;
}
return this.resolveAttribute(attrIdx.toNumber());
}
parseType(reader) {
// Parse a type reference from the reader (varint index)
const index = reader.parseVarInt().toNumber();
return this.resolveType(index);
}
resolveEntry(entries, index, entryType, depth = 0) {
// Simplified version of C++ resolveEntry - doesn't handle deferred worklists
// but does track recursion depth
const oldResolving = this._resolving;
this._resolving = true;
try {
const result = this.readEntry(entries, index, entryType, depth);
return result;
} finally {
this._resolving = oldResolving;
}
}
readEntry(entries, index, entryType, depth = 0) {
if (index >= entries.length) {
throw new mlir.Error(`Invalid '${entryType}' index '${index}'.`);
}
const entry = entries[index];
if (entry.resolved !== null) {
return entry.resolved;
}
if (depth > this.maxAttrTypeDepth) {
throw new mlir.Error(`Exceeded maximum '${entryType}' depth.`);
}
entry.resolved = { name: 'pending', value: `<${entryType} ${index}>` };
const entryData = this._sectionData.subarray(entry.offset, entry.offset + entry.size);
const reader = new _.EncodingReader(entryData);
const startPosition = reader.position;
if (entry.hasCustomEncoding) {
entry.resolved = this.parseCustomEntry(entry, reader, entryType, index, depth);
} else {
entry.resolved = this.parseAsmEntry(reader, entryType);
}
const bytesRead = reader.position - startPosition;
if (bytesRead > entry.size) {
throw new mlir.Error(`Read ${bytesRead} bytes but entry size is ${entry.size} bytes.`);
}
return entry.resolved;
}
parseCustomEntry(entry, reader, entryType, index, depth) {
// Lazy load the dialect interface (like BytecodeDialect::load)
const context = this.fileLoc.context;
if (entry.dialect.interface === undefined) {
entry.dialect.interface = context.getOrLoadDialect(entry.dialect.name);
}
if (!entry.dialect.interface) {
if (context.undocumented(entry.dialect.name)) {
throw new mlir.Error(`Undocumented dialect '${entry.dialect.name}'.`);
} else {
throw new mlir.Error(`Unsupported dialect '${entry.dialect.name}'.`);
}
}
const dialect = entry.dialect.interface;
const dialectReader = new _.DialectReader(this, this.stringReader, this.resourceReader, this.dialectsMap, reader, this.bytecodeVersion, depth);
switch (entryType) {
case 'type':
return dialect.readType(dialectReader);
case 'attr':
return dialect.readAttribute(dialectReader);
default:
throw new mlir.Error(`Unknown entry type '${entryType}'.`);
}
}
parseAsmEntry(reader, entryType) {
const asmStr = reader.parseNullTerminatedString();
const context = this.fileLoc.context;
if (entryType === 'type') {
const { type, numRead } = _.Parser.parseType(asmStr, context);
if (!type || numRead !== asmStr.length) {
throw new mlir.Error(`Failed to parse type '${asmStr}'.`);
}
return type;
}
const { attribute, numRead } = _.Parser.parseAttribute(asmStr, context, null);
if (!attribute || numRead !== asmStr.length) {
throw new mlir.Error(`Failed to parse attribute '${asmStr}'.`);
}
return attribute;
}
};
_.StringSectionReader = class {
initialize(sectionData) {
const decoder = new TextDecoder('utf-8');
const stringReader = new _.EncodingReader(sectionData);
const numStrings = stringReader.parseVarInt().toNumber();
this.strings = new Array(numStrings);
let stringDataEndOffset = sectionData.length;
for (let i = numStrings - 1; i >= 0; i--) {
const stringSize = stringReader.parseVarInt().toNumber();
if (stringDataEndOffset < stringSize) {
throw new mlir.Error('String size exceeds the available data size.');
}
const stringOffset = stringDataEndOffset - stringSize;
const buffer = sectionData.subarray(stringOffset, stringOffset + stringSize - 1);
this.strings[i] = decoder.decode(buffer);
stringDataEndOffset = stringOffset;
}
if ((sectionData.length - stringReader.size()) !== stringDataEndOffset) {
throw new mlir.Error('Unexpected trailing data between the offsets for strings and their data.');
}
}
parseString(reader) {
return reader.parseEntry(this.strings, 'string');
}
parseStringWithFlag(reader) {
const [entryIdx, flag] = reader.parseVarIntWithFlag();
const str = this.parseStringAtIndex(reader, entryIdx.toNumber());
return [str, flag];
}
parseStringAtIndex(reader, index) {
if (index >= this.strings.length) {
throw new mlir.Error('Invalid string index.');
}
return this.strings[index];
}
};
_.PropertiesSectionReader = class {
constructor() {
this._properties = new Map();
}
initialize(sectionData) {
const reader = new _.EncodingReader(sectionData);
const count = reader.parseVarInt().toNumber();
this._properties = new Array(count);
for (let i = 0; i < this._properties.length; i++) {
const size = reader.parseVarInt().toNumber();
const data = reader.parseBytes(size);
this._properties[i] = data;
}
}
read(fileLoc, dialectReader, opName, opState) {
const propIdx = dialectReader.readVarInt();
if (propIdx >= this._properties.length) {
return;
}
const propData = this._properties[propIdx];
if (propData.length === 0) {
return;
}
const propReader = dialectReader.withEncodingReader(new _.EncodingReader(propData));
// Reference: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
// Generated readProperties reads attributes in alphabetical order by name.
// Optional attributes (OptionalAttr, UnitAttr) use readOptionalAttribute (varint with flag).
const metadata = opName.metadata;
if (metadata && Array.isArray(metadata.attributes)) {
const sorted = metadata.attributes.slice().sort((a, b) => a.name.localeCompare(b.name));
for (const entry of sorted) {
if (propReader.reader.position >= propData.length) {
break;
}
const typeName = entry.type && entry.type.name ? entry.type.name : String(entry.type);
const isOptional = typeName === 'OptionalAttr' || typeName === 'UnitAttr' || typeName === 'UnitProp' || typeName === 'DefaultValuedOptionalAttr';
const attr = isOptional ? propReader.readOptionalAttribute() : propReader.readAttribute();
if (attr !== null && attr !== undefined) {
opState.addAttribute(entry.name, attr);
}
}
return;
}
// Fallback: read a single attribute as propertiesAttr
opState.propertiesAttr = propReader.readAttribute();
}
};
_.AsmResourceEntryKind = {
Blob: 0,
Bool: 1,
String: 2
};
_.ParsedResourceEntry.Bytecode = class {
constructor(key, kind, reader, stringReader) {
this.key = key;
this.kind = kind;
this.reader = reader;
this.stringReader = stringReader;
}
parseAsBlob() {
if (this.kind === _.AsmResourceEntryKind.Blob) {
const [data] = this.reader.parseBlobAndAlignment();
return data;
}
return null;
}
};
_.ResourceSectionReader = class {
constructor() {
this.dialectResources = [];
this.dialectResourceHandleRenamingMap = new Map();
}
initialize(fileLoc, config, dialects, stringReader, sectionData, offsetSectionData, dialectReader) {
const resourceReader = new _.EncodingReader(sectionData);
const offsetReader = new _.EncodingReader(offsetSectionData);
const numExternalResourceGroups = offsetReader.parseVarInt().toNumber();
const parseGroup = (handler, allowEmpty = false, processKeyFn = null) => {
const resolveKey = (key) => {
const remapped = this.dialectResourceHandleRenamingMap.get(key);
return remapped === undefined ? key : remapped;
};
return this.parseResourceGroup(fileLoc, allowEmpty, offsetReader, resourceReader, stringReader, handler, resolveKey, processKeyFn);
};
for (let i = 0; i < numExternalResourceGroups; i++) {
const key = stringReader.parseString(offsetReader);
const handler = config && config.getResourceParser ? config.getResourceParser(key) : null;
if (!parseGroup(handler)) {
return false;
}
}
const ctx = fileLoc.context;
while (!offsetReader.empty()) {
const dialect = offsetReader.parseEntry(dialects, "dialect");
dialect.load(dialectReader, ctx);
const handler = dialect.getLoadedDialect();
if (!handler) {
throw new mlir.Error(`Unknown dialect '${dialect.name}'.`);
}
const processResourceKeyFn = (key) => {
const handle = handler.declareResource(key);
if (!handle) {
throw new mlir.Error(`Unknown 'resource' key '${key}' for dialect '${dialect.name}'.`);
}
this.dialectResourceHandleRenamingMap.set(key, handler.getResourceKey(handle));
this.dialectResources.push(handler);
return true;
};
if (!parseGroup(handler, true, processResourceKeyFn)) {
return false;
}
}
return true;
}
parseResourceGroup(fileLoc, allowEmpty, offsetReader, resourceReader, stringReader, handler, remapKey, processKeyFn) {
const numResources = offsetReader.parseVarInt().toNumber();
for (let i = 0; i < numResources; i++) {
const key = stringReader.parseString(offsetReader);
const resourceOffset = offsetReader.parseVarInt().toNumber();
const kind = offsetReader.parseByte();
const data = resourceReader.parseBytes(resourceOffset);
if (processKeyFn && !processKeyFn(key)) {
return false;
}
if (allowEmpty && data.length === 0) {
continue;
}
const entryReader = new _.EncodingReader(data);
const resolvedKey = remapKey ? remapKey(key) : key;
const entry = new _.ParsedResourceEntry.Bytecode(resolvedKey, kind, entryReader, stringReader);
if (!handler) {
continue;
}
handler.parseResource(entry);
if (!entryReader.empty()) {
throw new mlir.Error(`Unexpected trailing bytes in resource entry '${resolvedKey}'.`);
}
}
return true;
}
parseResourceHandle(reader) {
return reader.parseEntry(this.dialectResources, "resource handle");
}
};
_.Location = class {
constructor(loc) {
this.impl = loc;
}
get context() {
return this.impl.context;
}
toString() {
return this.impl.toString();
}
};
_.LocationAttr = class {
constructor() {
this.name = 'loc';
}
printLocationInternal() {
throw new mlir.Error(`Unsupported location type '${this.constructor.name}'.`);
}
toString() {
return `loc(${this.printLocationInternal()})`;
}
};
_.UnknownLoc = class extends _.LocationAttr {
constructor(context) {
super();
this.context = context;
}
printLocationInternal() {
return 'unknown';
}
};
_.FileLineColLoc = class extends _.LocationAttr {
constructor(context, filename, line, col) {
super();
this.context = context;
this.filename = filename;
this.line = line;
this.col = col;
}
printLocationInternal() {
return `${this.line}:${this.col}`;
}
};
_.FileLineColRange = class extends _.LocationAttr {
constructor(filename, startLine, startCol, endLine, endCol) {
super();
this.filename = filename;
this.startLine = startLine;
this.startCol = startCol;
this.endLine = endLine;
this.endCol = endCol;
}
printLocationInternal() {
return `${this.filename}:${this.startLine}:${this.startCol} to ${this.endLine}:${this.endCol}`;
}
};
_.NameLoc = class extends _.LocationAttr {
constructor(name, childLoc) {
super();
this._name = name;
this.childLoc = childLoc;
}
printLocationInternal() {
if (this.childLoc) {
return `"${this._name}"(${this.childLoc.printLocationInternal()})`;
}
return `"${this._name}"`;
}
};
_.CallSiteLoc = class extends _.LocationAttr {
constructor(callee, caller) {
super();
this.callee = callee;
this.caller = caller;
}
printLocationInternal() {
return `callsite(${this.callee.printLocationInternal()} at ${this.caller.printLocationInternal()})`;
}
};
_.FusedLoc = class extends _.LocationAttr {
constructor(locations, metadata) {
super();
this.locations = locations;
this.metadata = metadata;
}
printLocationInternal() {
const locStrs = this.locations.map((loc) => loc.printLocationInternal());
if (this.metadata) {
return `fused<${this.metadata}>[${locStrs.join(', ')}]`;
}
return `fused[${locStrs.join(', ')}]`;
}
};
_.XtenNNDictLoc = class extends _.LocationAttr {
constructor(dictContent) {
super();
this.dictContent = dictContent;
}
printLocationInternal() {
return `#xten_nn<dict_loc(${this.dictContent})`;
}
};
_.OpaqueLoc = class extends _.LocationAttr {
constructor(index, identifier, fallbackLoc) {
super();
this.index = index;
this.identifier = identifier;
this.fallbackLoc = fallbackLoc || new _.UnknownLoc();
}
printLocationInternal() {
return this.fallbackLoc.printLocationInternal();
}
};
_.BytecodeReader = class {
constructor(reader, config) {
this.reader = new _.EncodingReader(reader);
this.config = config;
this.fileLoc = new _.FileLineColLoc(config.context, '', 0, 0);
this.valueScopes = [];
this.opNames = [];
this.dialects = [];
this.dialectsMap = new Map();
this.resourceReader = new _.ResourceSectionReader();
this.stringReader = new _.StringSectionReader();
this.propertiesReader = new _.PropertiesSectionReader();
this.version = { value: 0 };
this.attrTypeReader = new _.AttrTypeReader(this.stringReader, this.resourceReader, this.dialectsMap, this.version, this.fileLoc, config);
// Store buffer start address for alignment validation
this.bufferStart = reader instanceof Uint8Array ? reader.byteOffset : 0;
}
read() {
const reader = this.reader;
const signature = reader.parseBytes(4);
if (String.fromCharCode(...signature) !== 'ML\xEFR') {
throw new mlir.Error('Invalid MLIR bytecode signature.');
}
this.parseVersion(reader);
this.producer = reader.parseNullTerminatedString();
const sectionDatas = new Map();
while (reader.position < reader.length) {
const sectionIDAndHasAlignment = reader.parseByte();
const sectionID = sectionIDAndHasAlignment & 0x7F;
const length = reader.parseVarInt().toNumber();
const hasAlignment = sectionIDAndHasAlignment & 0x80;
if (sectionID >= 9) {
throw new mlir.Error(`Unsupported section identifier '${sectionID}'.`);
}
if (hasAlignment) {
const alignment = reader.parseVarInt().toNumber();
this.checkSectionAlignment(alignment);
reader.alignTo(alignment);
}
const sectionData = reader.parseBytes(length);
sectionDatas.set(sectionID, sectionData);
}
for (let sectionID = 0; sectionID < 9; sectionID++) {
if (sectionDatas.has(sectionID)) {
continue;
}
if (sectionID <= 4 || (sectionID === 8 && this.version.value >= 5)) { // kString, kDialect, kAttrType, kAttrTypeOffset, kIR, kProperties + kNativePropertiesEncoding
throw new mlir.Error(`Missing section '${sectionID}'.`);
}
}
this.stringReader.initialize(sectionDatas.get(0));
if (sectionDatas.has(8)) {
this.propertiesReader.initialize(sectionDatas.get(8));
}
this.parseDialectSection(sectionDatas.get(1));
this.parseResourceSection(sectionDatas.get(5), sectionDatas.get(6));
this.attrTypeReader.initialize(this.dialects, sectionDatas.get(2), sectionDatas.get(3));
return this.parseIRSection(sectionDatas.get(4));
}
checkSectionAlignment(alignment) {
// Check that the bytecode buffer meets the requested section alignment.
// In JavaScript, we validate the byteOffset within the ArrayBuffer.
// This ensures the buffer offset is aligned to the requested alignment.
if (!Number.isInteger(alignment) || alignment <= 0 || (alignment & (alignment - 1)) !== 0) {
throw new mlir.Error(`Invalid alignment value: ${alignment} (must be a power of 2).`);
}
const isGloballyAligned = (this.bufferStart & (alignment - 1)) === 0;
if (!isGloballyAligned) {
throw new mlir.Error(`Expected section alignment ${alignment} but bytecode buffer offset 0x${this.bufferStart.toString(16)} is not aligned.`);
}
}
parseVersion(reader) {
this.version.value = reader.parseVarInt().toNumber();
const kVersion = 6;
const kMinSupportedVersion = 0;
if (this.version.value < kMinSupportedVersion || this.version.value > kVersion) {
throw new mlir.Error(`Unsupported MLIR bytecode version '${this.version.value}'.`);
}
}
parseDialectSection(sectionData) {
const sectionReader = new _.EncodingReader(sectionData);
const numDialects = sectionReader.parseVarInt().toNumber();
const checkSectionAlignment = (alignment) => {
this.checkSectionAlignment(alignment);
};
for (let i = 0; i < numDialects; i++) {
this.dialects.push(new _.BytecodeDialect());
if (this.version.value < 1) { // kDialectVersioning
this.dialects[i].name = this.stringReader.parseString();
this.dialectsMap.set(this.dialects[i].name, this.dialects[i]);
continue;
}
const [dialectNameIdx, versionAvailable] = sectionReader.parseVarIntWithFlag();
this.dialects[i].name = this.stringReader.parseStringAtIndex(sectionReader, dialectNameIdx.toNumber());
if (versionAvailable) {
const [sectionID, versionBuffer] = sectionReader.parseSection(checkSectionAlignment);
this.dialects[i].versionBuffer = versionBuffer;
if (sectionID !== 7) { // kDialectVersion
throw new mlir.Error(`Expected dialect version section.`);
}
}
this.dialectsMap.set(this.dialects[i].name, this.dialects[i]);
}
let index = 0;
let numOps = -1;
const parseOpName = (dialect) => {
const opName = {};
opName.dialect = dialect;
if (this.version.value < 5) { // kNativePropertiesEncoding
opName.name = this.stringReader.parseString(sectionReader);
} else {
[opName.name, opName.wasRegistered] = this.stringReader.parseStringWithFlag(sectionReader);
}
if (numOps < 0) {
this.opNames.push(opName);
} else {
this.opNames[index++] = opName;
}
};
if (this.version.value >= 4) { // kElideUnknownBlockArgLocation
numOps = sectionReader.parseVarInt().toNumber();
this.opNames = new Array(numOps);
}
while (sectionReader.position < sectionData.length) {
_.BytecodeReader.parseDialectGrouping(sectionReader, this.dialects, parseOpName);
}
}
parseResourceSection(resourceData, resourceOffsetData) {
if (!resourceOffsetData) {
return true;
}
const dialectReader = new _.DialectReader(this.attrTypeReader, this.stringReader, this.resourceReader, this.dialectsMap, new _.EncodingReader(resourceData || new Uint8Array(0)), this.version);
return this.resourceReader.initialize(this.fileLoc, this.config, this.dialects, this.stringReader, resourceData, resourceOffsetData, dialectReader);
}
parseIRSection(sectionData) {
const reader = new _.EncodingReader(sectionData);
const block = { operations: [] };
this.valueScopes = [[]];
const regionStack = [{
block,
curRegion: 0,
numRegions: 1,
curBlock: 0,
numBlocks: 1,
numOpsRemaining: 0,
numValues: 0,
blocks: [block],
nextValueIdx: 0,
isTopLevel: true,
reader
}];
const firstBlockHeader = this.parseBlockHeader(reader);
regionStack[0].numOpsRemaining = firstBlockHeader.numOps;
if (firstBlockHeader.hasArgs) {
const [scope] = this.valueScopes;
this.parseBlockArguments(reader, block, scope, 0);
regionStack[0].nextValueIdx = block.arguments ? block.arguments.length : 0;
}
// Iteratively parse regions until everything has been resolved.
while (regionStack.length > 0) {
this.parseRegions(regionStack, regionStack[regionStack.length - 1]);
}
return block;
}
parseRegion(readState) {
const reader = readState.reader;
const numBlocks = reader.parseVarInt().toNumber();
if (numBlocks === 0) {
return false;
}
const numValues = reader.parseVarInt().toNumber();
readState.numValues = numValues;
readState.numBlocks = numBlocks;
// Create the blocks within this region.
const blocks = [];
for (let j = 0; j < numBlocks; j++) {
blocks.push({ operations: [], arguments: [] });
}
readState.blocks = blocks;
readState.region.blocks = blocks;
// Prepare the current value scope for this region.
const scope = this.valueScopes[this.valueScopes.length - 1];
const valueOffset = scope.length;
readState.valueOffset = valueOffset;
for (let j = 0; j < numValues; j++) {
scope.push(null);
}
// Parse the entry block of the region.
readState.curBlock = 0;
const blockHeader = this.parseBlockHeader(reader);
readState.numOpsRemaining = blockHeader.numOps;
if (blockHeader.hasArgs) {
this.parseBlockArguments(reader, blocks[0], scope, valueOffset);
}
const numBlockArgs = blocks[0].arguments ? blocks[0].arguments.length : 0;
readState.nextValueIdx = valueOffset + numBlockArgs;
return true;
}
parseRegions(regionStack, readState) {
const reader = readState.reader;
while (true) {
while (readState.numOpsRemaining > 0) {
readState.numOpsRemaining--;
// Read in the next operation. We don't read its regions directly,
// we handle those afterwards as necessary.
const { state: opState, resultNames, isIsolatedFromAbove, resultIndices } = this.parseOpWithoutRegions(reader, readState);
const op = _.Operation.create(opState);
// Assign result names for Netron display (reference: names are in parser symbol table)
// Also update the value scope to replace placeholders with actual OpResult objects
const scope = this.valueScopes[this.valueScopes.length - 1];
for (let i = 0; i < resultNames.length && i < op.results.length; i++) {
op.results[i].name = resultNames[i];
// Replace placeholder in scope with actual result
if (resultIndices && i < resultIndices.length) {
const valueIdx = resultIndices[i];
if (valueIdx < scope.length) {
scope[valueIdx] = op.results[i];
}
}
}
readState.blocks[readState.curBlock].operations.push(op);
if (op.regions && op.regions.length > 0) {
for (let i = op.regions.length - 1; i >= 0; i--) {
const region = op.regions[i];
const childState = {
region,
curRegion: 0,
numRegions: 1,
curBlock: 0,
numBlocks: 0,
numOpsRemaining: 0,
numValues: 0,
blocks: [],
valueOffset: 0,
nextValueIdx: 0,
isIsolated: isIsolatedFromAbove,
reader,
owningReader: null
};
// Isolated regions are encoded as a section in version 2 and above.
if (this.version.value >= 2 && isIsolatedFromAbove) { // kLazyLoading
const checkSectionAlignment = (alignment) => {
this.checkSectionAlignment(alignment);
};
const [sectionID, sectionData] = reader.parseSection(checkSectionAlignment);
if (sectionID !== 4) { // kIR
throw new mlir.Error(`Expected IR section for region.`);
}
childState.owningReader = new _.EncodingReader(sectionData);
childState.reader = childState.owningReader;
}
// If the op is isolated from above, push a new value scope.
if (isIsolatedFromAbove) {
this.valueScopes.push([]);
}
// Parse the region and push to stack if non-empty
if (this.parseRegion(childState)) {
regionStack.push(childState);
return;
}
}
}
}
// Move to the next block of the region.
readState.curBlock++;
if (readState.curBlock >= readState.numBlocks) {
break;
}
const blockHeader = this.parseBlockHeader(reader);
readState.numOpsRemaining = blockHeader.numOps;
if (blockHeader.hasArgs) {
const scope = this.valueScopes[this.valueScopes.length - 1];
const argOffset = readState.nextValueIdx === undefined ? 0 : readState.nextValueIdx;
this.parseBlockArguments(reader, readState.blocks[readState.curBlock], scope, argOffset);
const numBlockArgs = readState.blocks[readState.curBlock].arguments ? readState.blocks[readState.curBlock].arguments.length : 0;
if (readState.nextValueIdx !== undefined) {
readState.nextValueIdx += numBlockArgs;
}
}
}
if (readState.isIsolated) {
this.valueScopes.pop();
}
regionStack.pop();
}
parseBlockHeader(reader) {
const numOpsAndHasArgs = reader.parseVarInt();
const numOps = (numOpsAndHasArgs >> 1n).toNumber();
const hasArgs = (numOpsAndHasArgs & 1n) === 1n;
return { numOps, hasArgs };
}
parseBlockArguments(reader, block, scope, valueOffset) {
const numArgs = reader.parseVarInt().toNumber();
block.arguments = [];
for (let i = 0; i < numArgs; i++) {
let type = null;
let location = null;
if (this.version.value >= 4) { // kElideUnknownBlockArgLocation
const [typeIdx, hasLocation] = reader.parseVarIntWithFlag();
type = this.attrTypeReader.readType(typeIdx.toNumber());
if (hasLocation) {
const locIdx = reader.parseVarInt().toNumber();
location = this.attrTypeReader.readAttribute(locIdx);
}
} else {
const typeIdx = reader.parseVarInt().toNumber();
type = this.attrTypeReader.readType(typeIdx);
const locIdx = reader.parseVarInt().toNumber();
location = this.attrTypeReader.readAttribute(locIdx);
}
const argName = `%${valueOffset + i}`;
const arg = new _.Value(argName, type);
arg.location = location;
block.arguments.push(arg);
if (scope && (valueOffset + i) < scope.length) {
scope[valueOffset + i] = arg;
}
}
// Use-list ordering (version >= 3) - stored after all arguments
// hasUseListOrders byte is read whenever hasArgs is true (not dependent on numArgs)
if (this.version.value >= 3) { // kUseListOrdering
const hasUseListOrders = reader.parseByte();
if (hasUseListOrders !== 0) {
const orders = this.parseUseListOrdersForRange(reader, numArgs);
// Store ordering information on block arguments
for (const order of orders) {
if (order.argIndex < block.arguments.length) {
block.arguments[order.argIndex].useListOrder = order.ordering;
}
}
}
}
}
parseUseListOrdersForRange(reader, numValues) {
const orders = [];
let numToRead = 1;
if (numValues > 1) {
numToRead = reader.parseVarInt().toNumber();
}
for (let i = 0; i < numToRead; i++) {
let argIndex = 0;
if (numValues > 1) {
argIndex = reader.parseVarInt().toNumber();
}
const numUsesAndFlag = reader.parseVarInt();
const numUses = (numUsesAndFlag >> 1n).toNumber();
const indexPairEncoding = (numUsesAndFlag & 1n) === 1n;
const indices = [];
for (let j = 0; j < numUses; j++) {
indices.push(reader.parseVarInt().toNumber());
}
orders.push({ argIndex, ordering: { indexPairEncoding, indices } });
}
return orders;
}
parseOpName(reader) {
const opName = reader.parseEntry(this.opNames, 'operation name');
if (!opName.opName) {
const name = `${opName.dialect.name}.${opName.name}`;
opName.opName = _.RegisteredOperationName.lookup(name, this.fileLoc.context);
if (!opName.opName) {
throw new mlir.Error(`Unregistered bytecode operation '${name}'.`);
}
}
return [opName.opName, opName.wasRegistered];
}
parseOpWithoutRegions(reader, state) {
// Parse operation name index
const [opName, wasRegistered] = this.parseOpName(reader);
const opMask = reader.parseByte();
const kHasAttrs = 0x01;
const kHasResults = 0x02;
const kHasOperands = 0x04;
const kHasSuccessors = 0x08;
const kHasInlineRegions = 0x10;
const kHasUseListOrders = 0x20;
const kHasProperties = 0x40;
const opState = new _.OperationState(null, opName);
const locIdx = reader.parseVarInt().toNumber();
opState.location = this.attrTypeReader.readAttribute(locIdx);
if (opMask & kHasAttrs) {
const dictAttrIdx = reader.parseVarInt().toNumber();
const dictAttr = this.attrTypeReader.readAttribute(dictAttrIdx);
if (dictAttr && dictAttr.value) {
if (dictAttr.value instanceof Map) {
// Already parsed as Map from custom-encoded DictionaryAttr
opState.attributes = dictAttr.value;
} else if (typeof dictAttr.value === 'string') {
// Parse dictionary attribute from ASM string format
opState.attributes = this.parseAttributeDict(dictAttr.value);
}
}
}
// Parse properties (version >= 5)
if (opMask & kHasProperties) {
if (wasRegistered) {
const dialectReader = new _.DialectReader(this.attrTypeReader, this.stringReader, this.resourceReader, this.dialectsMap, reader, this.version);
this.propertiesReader.read(this.fileLoc, dialectReader, opName, opState);
} else {
// Unregistered operations store properties as a single dictionary attribute
const propAttrIdx = reader.parseVarInt().toNumber();
const propAttr = this.attrTypeReader.readAttribute(propAttrIdx);
if (propAttr && propAttr.value) {
const propAttrs = this.parseAttributeDict(propAttr.value);
for (const [key, value] of propAttrs) {
opState.addAttribute(key, value);
}
}
}
}
// Parse results - add types to OperationState.types, track values in scope
const resultNames = [];
const resultIndices = [];
if (opMask & kHasResults) {
const numResults = reader.parseVarInt().toNumber();
const scope = this.valueScopes[this.valueScopes.length - 1];
for (let i = 0; i < numResults; i++) {
const typeIdx = reader.parseVarInt().toNumber();
const type = this.attrTypeReader.readType(typeIdx);
opState.addTypes([type]);
const valueIdx = state && state.nextValueIdx !== undefined ? state.nextValueIdx++ : scope.length;
const valueName = `%${valueIdx}`;
resultNames.push(valueName);
resultIndices.push(valueIdx);
// Create placeholder in scope (will be replaced after Operation creation with actual OpResult)
const placeholder = new _.Value(valueName, type);
if (valueIdx < scope.length) {
scope[valueIdx] = placeholder;
} else {
scope.push(placeholder);
}
}
}
if (opMask & kHasOperands) {
const numOperands = reader.parseVarInt().toNumber();
for (let i = 0; i < numOperands; i++) {
const valueIdx = reader.parseVarInt().toNumber();
const scope = this.valueScopes[this.valueScopes.length - 1];
if (valueIdx < scope.length && scope[valueIdx]) {
opState.operands.push(scope[valueIdx]);
} else {
opState.operands.push(new _.Value(`%${valueIdx}`, null));
}
}
}
if (opMask & kHasSuccessors) {
const numSuccessors = reader.parseVarInt().toNumber();
opState.successors = [];
for (let i = 0; i < numSuccessors; i++) {
const blockIdx = reader.parseVarInt().toNumber();
opState.successors.push(blockIdx);
}
}
// Parse use-list orders (version >= 3)
if (this.version.value >= 3 && (opMask & kHasUseListOrders)) { // kUseListOrdering
const numResults = opState.types.length;
const orders = this.parseUseListOrdersForRange(reader, numResults);
if (orders.length > 0) {
opState.useListOrders = orders;
}
}
let isIsolatedFromAbove = false;
if (opMask & kHasInlineRegions) {
const numRegionsAndIsIsolated = reader.parseVarInt();
const numRegions = (numRegionsAndIsIsolated >> 1n).toNumber();
isIsolatedFromAbove = (numRegionsAndIsIsolated & 1n) === 1n;
for (let i = 0; i < numRegions; i++) {
opState.regions.push({ blocks: [] });
}
}
return { state: opState, resultNames, isIsolatedFromAbove, resultIndices };
}
parseAttributeDict(str) {
const attrs = new Map();
// Parse dictionary attribute format: {key = value, ...}
if (!str.startsWith('{') || !str.endsWith('}')) {
return attrs;
}
const content = str.slice(1, -1).trim();
if (content.length === 0) {
return attrs;
}
let i = 0;
while (i < content.length) {
while (i < content.length && /\s/.test(content[i])) {
i++;
}
if (i >= content.length) {
break;
}
// Read key (alphanumeric/underscore)
const keyStart = i;
while (i < content.length && /[a-zA-Z0-9_]/.test(content[i])) {
i++;
}
const key = content.slice(keyStart, i);
if (!key) {
break;
}
while (i < content.length && /\s/.test(content[i])) {
i++;
}
if (i >= content.length || content[i] !== '=') {
break;
}
i++; // skip '='
while (i < content.length && /\s/.test(content[i])) {
i++;
}
// Read value (until balanced comma or end)
const valueStart = i;
let depth = 0;
let inString = false;
while (i < content.length) {
const c = content[i];
if (inString) {
if (c === '"' && content[i - 1] !== '\\') {
inString = false;
}
i++;
continue;
}
if (c === '"') {
inString = true;
i++;
continue;
}
if (c === '{' || c === '[' || c === '(' || c === '<') {
depth++;
} else if (c === '}' || c === ']' || c === ')' || c === '>') {
depth--;
} else if (c === ',' && depth === 0) {
break;
}
i++;
}
const value = content.slice(valueStart, i).trim();
attrs.set(key, { value, toString: () => value });
if (i < content.length && content[i] === ',') {
i++;
}
}
return attrs;
}
static parseDialectGrouping(reader, dialects, entryCallback) {
const dialect = reader.parseEntry(dialects, 'dialect');
const numEntries = reader.parseVarInt().toNumber();
for (let i = 0; i < numEntries; i++) {
entryCallback(dialect);
}
}
};
_.Constraint = class {
constructor(name, args, values) {
this.name = name;
this.args = args || null;
this.values = values || null;
}
static parse(value) {
const tokens = _.Constraint.tokenize(value);
const result = _.Constraint.parseTokens(tokens, 0);
const { name, args, values } = result.value;
return new _.Constraint(name, args, values);
}
toString() {
if (this.args && this.args.length > 0) {
const args = this.args.map((arg) => arg.name).join(', ');
return `${this.name}<${args}>`;
}
if (this.values && this.values.length > 0) {
return `${this.name}{${this.values.join(' | ')}}`;
}
return this.name;
}
static tokenize(str) {
const tokens = [];
let i = 0;
while (i < str.length) {
const ch = str[i];
if (/\s/.test(ch)) {
i++;
continue;
}
if ('<>={}[](),|'.indexOf(ch) !== -1) {
tokens.push({ type: ch, value: ch, pos: i });
i++;
continue;
}
if (ch === ':' && i + 1 < str.length && str[i + 1] === ':') {
tokens.push({ type: '::', value: '::', pos: i });
i += 2;
continue;
}
if (ch === ':') {
i++;
continue;
}
if (ch === '"' || ch === "'") {
const quote = ch;
let j = i + 1;
while (j < str.length && str[j] !== quote) {
if (str[j] === '\\' && j + 1 < str.length) {
j += 2;
} else {
j++;
}
}
if (j < str.length) {
tokens.push({ type: 'string', value: str.substring(i + 1, j), pos: i });
i = j + 1;
} else {
tokens.push({ type: 'ident', value: str.substring(i), pos: i });
break;
}
continue;
}
if (/[a-zA-Z_0-9-]/.test(ch)) {
let j = i;
while (j < str.length && /[a-zA-Z_0-9-]/.test(str[j])) {
j++;
}
const ident = str.substring(i, j);
tokens.push({ type: 'ident', value: ident, pos: i });
i = j;
continue;
}
i++;
}
return tokens;
}
static parseTokens(tokens, pos) {
if (pos >= tokens.length) {
return null;
}
const token = tokens[pos];
if (token.type === '::') {
let name = '';
let nextPos = pos;
while (nextPos < tokens.length) {
if (tokens[nextPos].type === '::') {
name += '::';
nextPos++;
} else if (tokens[nextPos].type === 'ident') {
name += tokens[nextPos].value;
nextPos++;
} else {
break;
}
}
if (!name) {
return null;
}
if (nextPos < tokens.length) {
const nextToken = tokens[nextPos];
if (nextToken.type === '{') {
return _.Constraint._parseEnum(tokens, pos, name);
}
if (nextToken.type === '<') {
return _.Constraint._parseGeneric(tokens, pos, name);
}
}
return { value: { name }, nextPos };
}
if (token.type !== 'ident') {
return null;
}
let name = token.value;
let nextPos = pos + 1;
while (nextPos < tokens.length && tokens[nextPos].type === '::') {
nextPos++;
if (nextPos < tokens.length && tokens[nextPos].type === 'ident') {
const value = tokens[nextPos++].value;
name += `::${value}`;
} else {
break;
}
}
if (nextPos >= tokens.length) {
return { value: { name }, nextPos };
}
const nextToken = tokens[nextPos];
if (nextToken.type === '{') {
return _.Constraint._parseEnum(tokens, pos, name);
}
if (nextToken.type === '<') {
return _.Constraint._parseGeneric(tokens, pos, name);
}
return { value: { name }, nextPos };
}
static _parseEnum(tokens, startPos, name) {
let pos = startPos;
while (pos < tokens.length && (tokens[pos].type === 'ident' || tokens[pos].type === '::')) {
pos++;
}
if (pos >= tokens.length || tokens[pos].type !== '{') {
return null;
}
pos++;
const values = [];
let currentValue = '';
while (pos < tokens.length && tokens[pos].type !== '}') {
const token = tokens[pos];
if (token.type === '|') {
if (currentValue.trim()) {
values.push(currentValue.trim());
currentValue = '';
}
pos++;
} else if (token.type === 'ident') {
if (currentValue) {
currentValue += ' ';
}
currentValue += token.value;
pos++;
} else if (token.type === '::') {
currentValue += '::';
pos++;
} else {
pos++;
}
}
if (currentValue.trim()) {
values.push(currentValue.trim());
}
if (pos < tokens.length && tokens[pos].type === '}') {
pos++;
}
return { value: { name, values }, nextPos: pos };
}
static _parseGeneric(tokens, startPos, name) {
let pos = startPos;
while (pos < tokens.length && (tokens[pos].type === 'ident' || tokens[pos].type === '::')) {
pos++;
}
if (pos >= tokens.length || tokens[pos].type !== '<') {
return null;
}
pos++;
const args = [];
let angleDepth = 1;
let bracketDepth = 0;
let currentArg = [];
while (pos < tokens.length && (angleDepth > 0 || bracketDepth > 0)) {
const token = tokens[pos];
if (token.type === '<') {
angleDepth++;
currentArg.push(token);
pos++;
} else if (token.type === '>') {
angleDepth--;
if (angleDepth === 0 && bracketDepth === 0) {
if (currentArg.length > 0) {
const parsed = _.Constraint.parseArgumentTokens(currentArg);
if (parsed !== null) {
args.push(parsed);
}
}
pos++;
break;
}
currentArg.push(token);
pos++;
} else if (token.type === '[') {
bracketDepth++;
currentArg.push(token);
pos++;
} else if (token.type === ']') {
bracketDepth--;
currentArg.push(token);
pos++;
} else if (token.type === ',' && angleDepth === 1 && bracketDepth === 0) {
if (currentArg.length > 0) {
const parsed = _.Constraint.parseArgumentTokens(currentArg);
if (parsed !== null) {
args.push(parsed);
}
currentArg = [];
}
pos++;
} else {
currentArg.push(token);
pos++;
}
}
return { value: { name, args }, nextPos: pos };
}
static parseArgumentTokens(tokens) {
if (!tokens || tokens.length === 0) {
return null;
}
tokens = tokens.filter((t) => t.type !== undefined);
if (tokens[0].type === '[') {
return _.Constraint.parseListArgument(tokens);
}
if (tokens[0].type === 'string') {
return tokens[0].value;
}
if (tokens[0].type === 'ident' || tokens[0].type === '::') {
const result = _.Constraint.parseTokens(tokens, 0);
if (result && result.nextPos === tokens.length) {
return result.value;
}
}
let literal = '';
for (const token of tokens) {
if (token.type === 'ident' || token.type === 'string') {
if (literal && !/^[,[\]():.]$/.test(literal[literal.length - 1])) {
literal += ' ';
}
literal += token.value;
} else if (token.type === '::') {
literal += '::';
} else if ('{}[](),.'.indexOf(token.type) !== -1) {
literal += token.value;
}
}
return literal.trim() || null;
}
static parseListArgument(tokens) {
if (!tokens || tokens.length === 0 || tokens[0].type !== '[') {
return null;
}
let pos = 1;
const items = [];
let bracketDepth = 1;
let angleDepth = 0;
let currentItem = [];
while (pos < tokens.length && (bracketDepth > 0 || angleDepth > 0)) {
const token = tokens[pos];
if (token.type === '[') {
bracketDepth++;
currentItem.push(token);
pos++;
} else if (token.type === ']') {
bracketDepth--;
if (bracketDepth === 0 && angleDepth === 0) {
if (currentItem.length > 0) {
const parsed = _.Constraint.parseArgumentTokens(currentItem);
if (parsed !== null) {
items.push(parsed);
}
}
break;
}
currentItem.push(token);
pos++;
} else if (token.type === '<') {
angleDepth++;
currentItem.push(token);
pos++;
} else if (token.type === '>') {
angleDepth--;
currentItem.push(token);
pos++;
} else if (token.type === ',' && bracketDepth === 1 && angleDepth === 0) {
if (currentItem.length > 0) {
const parsed = _.Constraint.parseArgumentTokens(currentItem);
if (parsed !== null) {
items.push(parsed);
}
currentItem = [];
}
pos++;
} else {
currentItem.push(token);
pos++;
}
}
return items;
}
};
_.AssemblyFormatParser = class {
constructor(metadata) {
this._metadata = metadata;
this._buffer = metadata.assemblyFormat || '';
this._pos = 0;
}
match(char) {
return this._pos < this._buffer.length && this._buffer[this._pos] === char;
}
accept(str) {
if (str.length === 1) {
if (this.match(str)) {
this._pos++;
return true;
}
return false;
}
const remaining = this._buffer.substring(this._pos);
if (remaining.startsWith(str)) {
// Check that keyword is not followed by alphanumeric (to avoid "type" in "typename")
const nextChar = this._buffer[this._pos + str.length];
if (nextChar && /[a-zA-Z0-9_-]/.test(nextChar)) {
return false;
}
this._pos += str.length;
return true;
}
return false;
}
expect(char) {
if (!this.match(char)) {
throw new mlir.Error(`Expected '${char}'.`);
}
this._pos++;
}
parse() {
const directives = [];
this._skipWhitespace();
while (this._pos < this._buffer.length) {
const directive = this.parseDirective();
directives.push(directive);
this._skipWhitespace();
}
return directives;
}
parseDirective() {
const ch = this._buffer[this._pos];
if (!ch || this._pos >= this._buffer.length) {
throw new mlir.Error(`Unexpected end of format string.`);
}
// Parenthesized group: can be optional (...)? or conditional (...):(...) or just grouping (...)
if (this.match('(')) {
this.accept('(');
const elements = [];
let anchorElement = null;
this._skipWhitespace();
while (!this.match(')')) {
const elem = this.parseDirective();
if (elem.type === 'anchor') {
// Standalone anchor - applies to the previous element
if (elements.length > 0) {
const prev = elements[elements.length - 1];
anchorElement = prev.name || prev.type;
}
} else {
if (elem.anchor) {
anchorElement = elem.name || elem.type;
}
elements.push(elem);
}
this._skipWhitespace();
}
this.expect(')');
this._skipWhitespace();
// Check what follows to determine the group type
if (this.accept('?')) {
// Optional group: (...)?
return { type: 'optional_group', elements, anchor: anchorElement };
}
if (this.accept(':')) {
// Conditional alternative: (...):(...)?
this._skipWhitespace();
const secondAlt = [];
let isSecondOptional = false;
if (this.accept('(')) {
this._skipWhitespace();
while (!this.match(')')) {
const elem = this.parseDirective();
secondAlt.push(elem);
this._skipWhitespace();
}
this.expect(')');
this._skipWhitespace();
if (this.accept('?')) {
isSecondOptional = true;
}
}
return { type: 'conditional_alternative', firstAlt: elements, secondAlt, secondOptional: isSecondOptional };
}
return { type: 'group', elements };
}
// Literal: `keyword`
if (this.accept('`')) {
const value = this.parseUntil('`');
this.expect('`');
// MLIR reference: Empty literals (`` or ` `) are whitespace, not literals
if (value.length === 0 || value === ' ' || value === '\\n') {
return { type: 'whitespace', value }; // Return whitespace as a directive
}
return { type: 'literal', value };
}
if (this.accept('$')) {
const name = this.parseIdentifier();
const anchor = this.accept('^');
const metadata = this._metadata;
// Determine variable type from metadata first - matches reference implementation
// Check each metadata category in priority order
if (metadata.successors && metadata.successors.some((a) => a.name === name)) {
return { type: 'successor_ref', name, anchor };
}
if (metadata.attributes && metadata.attributes.some((a) => a.name === name)) {
return { type: 'attribute_ref', name, anchor };
}
if (metadata.operands && metadata.operands.some((a) => a.name === name)) {
return { type: 'operand_ref', name, anchor };
}
if (metadata.regions && metadata.regions.some((a) => a.name === name)) {
return { type: 'region_ref', name, anchor };
}
throw new mlir.Error(`Unknown variable '$${name}' in assembly format.`);
}
if (this.accept('type')) {
const args = this.parseParenList();
const anchor = this.accept('^');
return { type: 'type', args, anchor };
}
if (this.accept('qualified')) {
const args = this.parseParenList();
const anchor = this.accept('^');
return { type: 'qualified', args, anchor };
}
if (this.accept('attr-dict-with-keyword')) {
return { type: 'attr_dict_with_keyword' };
}
if (this.accept('attr-dict')) {
return { type: 'attr_dict' };
}
if (this.accept('prop-dict')) {
return { type: 'prop_dict' };
}
if (this.accept('functional-type')) {
const args = this.parseParenList();
const anchor = this.accept('^');
return { type: 'functional_type', args, anchor };
}
if (this.accept('params')) {
return { type: 'params' };
}
if (this.accept('struct')) {
this.expect('(');
const args = [];
while (!this.match(')')) {
this._skipWhitespace();
if (this.match(')')) {
break;
}
const arg = this.parseDirective();
args.push(arg);
this._skipWhitespace();
this.accept(',');
}
this.expect(')');
return { type: 'struct', args };
}
if (this.accept('ref')) {
this.expect('(');
const arg = this.parseDirective();
this._skipWhitespace();
this.expect(')');
return { type: 'ref', arg };
}
if (this.accept('custom')) {
this.expect('<');
const parser = this.parseUntil('>');
this.expect('>');
const args = this.parseParenList();
const anchor = this.accept('^');
return { type: 'custom', parser, args, anchor };
}
if (this.accept('oilist')) {
this._skipWhitespace();
this.expect('(');
let content = '';
let depth = 1;
while (this._pos < this._buffer.length && depth > 0) {
const ch = this._buffer[this._pos];
if (ch === '(') {
depth++;
content += ch;
this._pos++;
} else if (ch === ')') {
depth--;
if (depth > 0) {
content += ch;
}
this._pos++;
} else {
content += ch;
this._pos++;
}
}
return { type: 'oilist', content };
}
if (this.accept('operands')) {
return { type: 'operands' };
}
if (this.accept('results')) {
return { type: 'results' };
}
if (this.accept('regions')) {
return { type: 'regions' };
}
if (this.accept('successors')) {
return { type: 'successors' };
}
if (ch === '^') {
this._pos++;
return { type: 'anchor' };
}
if (/^[:()[\]{}<>,=|]/.test(ch)) {
this._pos++;
return { type: 'literal', value: ch };
}
const context = this._buffer.substring(Math.max(0, this._pos - 10), Math.min(this._buffer.length, this._pos + 10));
throw new mlir.Error(`Unexpected '${ch}' in assembly format '${context}...'.`);
}
parseIdentifier() {
let name = '';
while (this._pos < this._buffer.length) {
const ch = this._buffer[this._pos];
if (/[a-zA-Z0-9_]/.test(ch)) {
name += ch;
this._pos++;
} else {
break;
}
}
return name;
}
parseUntil(terminator) {
let value = '';
while (this._pos < this._buffer.length && this._buffer[this._pos] !== terminator) {
value += this._buffer[this._pos];
this._pos++;
}
return value;
}
parseParenList() {
this._skipWhitespace();
if (!this.accept('(')) {
return [];
}
this._skipWhitespace();
if (this.accept(')')) {
return [];
}
const items = [];
const parseElement = () => {
let element = '';
let depth = 0;
while (this._pos < this._buffer.length) {
this._skipWhitespace();
if (this.accept('"')) {
// String literal - consume as a unit
element += '"';
element += this.parseUntil('"');
element += '"';
this.expect('"');
} else if (this.accept('$')) {
element += '$';
const id = this.parseIdentifier();
element += id;
} else if (this.accept('(')) {
// Nested parentheses - include in element (e.g., type($list))
element += '(';
depth++;
} else if (this.match(')')) {
if (depth === 0) {
// End of this element
break;
}
element += ')';
this.accept(')');
depth--;
} else if (this.match(',') && depth === 0) {
// Comma at top level - end of this element
break;
} else if (this.match('-')) {
// Handle hyphenated identifiers like attr-dict, functional-type
element += '-';
this.accept('-');
} else {
// Plain identifier (e.g., "type" in type($list))
const id = this.parseIdentifier();
if (!id) {
throw new mlir.Error(`Unexpected '${this._buffer[this._pos]}' in assembly format directive list.`);
}
element += id;
}
}
return element.trim();
};
const first = parseElement();
if (!first) {
throw new mlir.Error('Expected element.');
}
items.push(first);
this._skipWhitespace();
while (this.accept(',')) {
this._skipWhitespace();
const elem = parseElement();
if (!elem) {
throw new mlir.Error('Expected element after comma');
}
items.push(elem);
this._skipWhitespace();
}
this.expect(')');
return items;
}
_skipWhitespace() {
while (this._pos < this._buffer.length && /\s/.test(this._buffer[this._pos])) {
this._pos++;
}
}
};
_.DialectContext = class {
constructor(metadata) {
const operations = metadata.operations;
this._dialects = new Map();
this._dialects.set('builtin', new _.BuiltinDialect(operations));
this._dialects.set('std', new _.Dialect(operations, 'std'));
this._dialects.set('bufferization', new _.BufferizationDialect(operations));
this._dialects.set('stablehlo', new _.stablehlo.StableHLODialect(operations));
this._dialects.set('vhlo', new _.vhlo.VhloDialect(operations));
this._dialects.set('interpreter', new _.InterpreterDialect(operations));
this._dialects.set('affine', new _.AffineDialect(operations));
this._dialects.set('asuka', new _.AsukaDialect(operations));
this._dialects.set('arith', new _.ArithDialect(operations));
this._dialects.set('async', new _.async.AsyncDialect(operations));
this._dialects.set('cf', new _.CFDialect(operations));
this._dialects.set('emitc', new _.emitc.EmitCDialect(operations));
this._dialects.set('complex', new _.Dialect(operations, 'complex'));
this._dialects.set('index', new _.Dialect(operations, 'index'));
this._dialects.set('pdl', new _.pdl.PDLDialect(operations));
this._dialects.set('ptr', new _.ptr.PtrDialect(operations));
this._dialects.set('ub', new _.Dialect(operations, 'ub'));
this._dialects.set('amdgpu', new _.AMDGPUDialect(operations));
this._dialects.set('nvgpu', new _.nvgpu.NVGPUDialect(operations));
this._dialects.set('nvvm', new _.NVVMDialect(operations));
this._dialects.set('rocdl', new _.ROCDLDialect(operations));
this._dialects.set('nvws', new _.triton.nvws.NVWSDialect(operations));
this._dialects.set('tti', new _.Dialect(operations, 'tti'));
this._dialects.set('omp', new _.OpenMPDialect(operations));
this._dialects.set('proton', new _.triton.proton.ProtonDialect(operations));
this._dialects.set('proton_gpu', new _.Dialect(operations, 'proton_gpu'));
this._dialects.set('arm_sme', new _.ArmSMEDialect(operations));
this._dialects.set('arm_neon', new _.ArmNeonDialect(operations));
this._dialects.set('arm_sve', new _.ArmSVEDialect(operations));
this._dialects.set('shard', new _.ShardDialect(operations));
this._dialects.set('smt', new _.smt.SMTDialect(operations));
this._dialects.set('lagrad', new _.Dialect(operations, 'lagrad'));
this._dialects.set('iree_codegen', new _.IREECodegenDialect(operations));
this._dialects.set('iree_encoding', new _.Dialect(operations, 'iree_encoding'));
this._dialects.set('test', new _.test.TestDialect(operations));
this._dialects.set('scf', new _.SCFDialect(operations));
this._dialects.set('shape', new _.ShapeDialect(operations));
this._dialects.set('sparse_tensor', new _.SparseTensorDialect(operations));
this._dialects.set('func', new _.FuncDialect(operations));
this._dialects.set('gpu', new _.GpuDialect(operations));
this._dialects.set('llvm', new _.llvm.LLVMDialect(operations));
this._dialects.set('xegpu', new _.xegpu.XeGPUDialect(operations));
this._dialects.set('memref', new _.MemRefDialect(operations));
this._dialects.set('vector', new _.VectorDialect(operations));
this._dialects.set('x86', new _.Dialect(operations, 'x86'));
this._dialects.set('x86vector', new _.Dialect(operations, 'x86vector'));
this._dialects.set('onnx', new _.ONNXDialect(operations));
this._dialects.set('krnl', new _.KrnlDialect(operations));
this._dialects.set('torch', new _.torch.TorchDialect(operations));
this._dialects.set('torch_c', new _.Dialect(operations, 'torch_c'));
this._dialects.set('hal', new _.HALDialect(operations));
this._dialects.set('hal_loader', new _.HALLoaderDialect(operations));
this._dialects.set('hal_inline', new _.Dialect(operations, 'hal_inline'));
this._dialects.set('util', new _.UtilDialect(operations));
this._dialects.set('mhlo', new _.hlo.MhloDialect(operations));
this._dialects.set('chlo', new _.hlo.ChloDialect(operations));
this._dialects.set('thlo', new _.THLODialect(operations));
this._dialects.set('flow', new _.FlowDialect(operations));
this._dialects.set('stream', new _.StreamDialect(operations));
this._dialects.set('iree_vector_ext', new _.IREEVectorExtDialect(operations));
this._dialects.set('iree_tensor_ext', new _.IREETensorExtDialect(operations));
this._dialects.set('linalg', new _.LinalgDialect(operations));
this._dialects.set('iree_linalg_ext', new _.Dialect(operations, 'iree_linalg_ext'));
this._dialects.set('linalg_ext', this._dialects.get('iree_linalg_ext'));
this._dialects.set('quant', new _.quant.QuantDialect(operations));
this._dialects.set('tensor', new _.TensorDialect(operations));
this._dialects.set('tosa', new _.TosaDialect(operations));
this._dialects.set('tf', new _.TFDialect(operations));
this._dialects.set('tf_saved_model', new _.Dialect(operations, 'tf_saved_model'));
this._dialects.set('tf_type', new _.TFTypeDialect(operations));
this._dialects.set('tf_device', new _.TFDeviceDialect(operations));
this._dialects.set('tf_executor', new _.TFExecutorDialect(operations));
this._dialects.set('tf_framework', new _.TFFrameworkDialect(operations));
this._dialects.set('tfr', new _.TFRDialect(operations));
this._dialects.set('corert', new _.CoreRTDialect(operations));
this._dialects.set('tfrt', new _.TFRTDialect(operations));
this._dialects.set('tfrt_fallback', new _.Dialect(operations, 'tfrt_fallback'));
this._dialects.set('tfrt_fallback_async', new _.TFRTFallbackAsyncDialect(operations));
this._dialects.set('tfl', new _.TFLDialect(operations));
this._dialects.set('stdx', new _.StdxDialect(operations));
this._dialects.set('vm', new _.vm.VMDialect(operations));
this._dialects.set('math', new _.MathDialect(operations));
this._dialects.set('tm_tensor', new _.torch.TMTensorDialect(operations));
this._dialects.set('ml_program', new _.MLProgramDialect(operations));
this._dialects.set('iree_gpu', new _.IREEGPUDialect(operations));
this._dialects.set('tile', new _.TileDialect(operations));
this._dialects.set('pxa', new _.PXADialect(operations));
this._dialects.set('irdl', new _.IRDLDialect(operations));
this._dialects.set('transform', new _.TransformDialect(operations));
this._dialects.set('wasmssa', new _.WasmSSADialect(operations));
this._dialects.set('spirv', new _.spirv.SPIRVDialect(operations));
this._dialects.set('spv', this._dialects.get('spirv'));
this._dialects.set('toy', new _.hlo.ToyDialect(operations));
this._dialects.set('top', new _.Dialect(operations, 'top'));
this._dialects.set('tpu', new _.Dialect(operations, 'tpu'));
this._dialects.set('sdfg', new _.SdfgDialect(operations));
this._dialects.set('sdir', this._dialects.get('sdfg'));
this._dialects.set('check', new _.Dialect(operations, 'check'));
this._dialects.set('tt', new _.triton.TritonDialect(operations));
this._dialects.set('ttg', new _.triton.gpu.TritonGPUDialect(operations));
this._dialects.set('triton_gpu', this._dialects.get('ttg'));
this._dialects.set('gluon', new _.triton.gluon.GluonDialect(operations));
this._dialects.set('ttng', new _.triton.nvidia_gpu.TritonNvidiaGPUDialect(operations));
this._dialects.set('nvidia_gpu', this._dialects.get('ttng'));
this._dialects.set('amdg', new _.triton.amdgpu.TritonAMDGPUDialect(operations));
this._dialects.set('amd_gpu', this._dialects.get('amdg'));
this._dialects.set('michelson', new _.MichelsonDialect(operations));
this._dialects.set('tensorrt', new _.tensorrt.TensorRTDialect(operations));
this._dialects.set('executor', new _.executor.ExecutorDialect(operations));
this._dialects.set('exec', this._dialects.get('executor'));
this._dialects.set('tfrt_test', new _.TFRTTestDialect(operations));
this._dialects.set('xevm', new _.XeVMDialect(operations));
this._dialects.set('vmvx', new _.VMVXDialect(operations));
this._dialects.set('mlrt', new _.MLRTDialect(operations));
this._dialects.set('tfrt_tensor', new _.TFRTTensorDialect(operations));
this._dialects.set('tfrt_dht', new _.TFRTDHTDialect(operations));
this._dialects.set('coo', new _.Dialect(operations, 'coo'));
this._dialects.set('tfd', new _.TFDDialect(operations));
this._dialects.set('acc', new _.ACCDialect(operations));
this._dialects.set('cuda', new _.Dialect(operations, 'cuda'));
this._dialects.set('trtrt', new _.Dialect(operations, 'trtrt'));
this._dialects.set('plan', new _.PlanDialect(operations));
this._dialects.set('kernel', new _.KernelDialect(operations));
this._dialects.set('nvg', new _.Dialect(operations, 'nvg'));
this._dialects.set('mpi', new _.Dialect(operations, 'mpi'));
this._dialects.set('pdl_interp', new _.PDLInterpDialect(operations));
this._dialects.set('standalone', new _.Dialect(operations, 'standalone'));
this._dialects.set('custom', new _.Dialect(operations, 'custom'));
this._dialects.set('layer', new _.Dialect(operations, 'layer'));
this._dialects.set('foo', new _.Dialect(operations, 'foo'));
this._dialects.set('some', new _.Dialect(operations, 'some'));
this._dialects.set('ts', new _.Dialect(operations, 'ts'));
this._dialects.set('tf_mlrt', new _.Dialect(operations, 'tf_mlrt'));
this._dialects.set('io_parameters', new _.IOParametersDialect(operations));
this._dialects.set('pcf', new _.PCFDialect(operations));
this._dialects.set('linalgx', new _.Dialect(operations, 'linalgx'));
this._dialects.set('xsmm', new _.XSMMDialect(operations));
this._dialects.set('sdy', new _.SdyDialect(operations));
this._dialects.set('mpmd', new _.MPMDDialect(operations));
this._dialects.set('tfg', new _.TFGDialect(operations));
this._dialects.set('vt', new _.Dialect(operations, 'vt'));
this._dialects.set('testd', new _.Dialect(operations, 'testd'));
this._dialects.set('cmath', new _.Dialect(operations, 'cmath'));
this._dialects.set('bytecode', new _.Dialect(operations, 'bytecode'));
this._dialects.set('test_irdl_to_cpp', new _.Dialect(operations, 'test_irdl_to_cpp'));
this._dialects.set('iree_unregistered', new _.Dialect(operations, 'iree_unregistered'));
this._dialects.set('cir', new _.Dialect(operations, 'cir'));
this._dialects.set('migraphx', new _.Dialect(operations, 'migraphx'));
this._dialects.set('xla', new _.XlaDialect(operations));
this._dialects.set('xla_gpu', new _.XlaGpuDialect(operations));
this._dialects.set('xla_cpu', new _.Dialect(operations, 'xla_cpu'));
this._dialects.set('xla_framework', new _.Dialect(operations, 'xla_framework'));
this._dialects.set('ifrt', new _.Dialect(operations, 'ifrt'));
this._dialects.set('vifrt', new _.Dialect(operations, 'vifrt'));
this._dialects.set('triton_xla', new _.triton.xla.TritonXlaDialect(operations));
this._dialects.set('xtile', new _.XTileDialect(operations));
this._dialects.set('xten_nn', new _.XtenNNDialect(operations));
this._dialects.set('ensemble', new _.EnsembleDialect(operations));
this._dialects.set('poly', new _.PolyDialect(operations));
this._dialects.set('noisy', new _.NoisyDialect(operations));
this._dialects.set('plugin', new _.Dialect(operations, 'plugin'));
this._dialects.set('ttcore', new _.Dialect(operations, 'ttcore'));
this._dialects.set('ttir', new _.Dialect(operations, 'ttir'));
this._dialects.set('ttnn', new _.Dialect(operations, 'ttnn'));
this._dialects.set('ttkernel', new _.Dialect(operations, 'ttkernel'));
this._dialects.set('ttmetal', new _.Dialect(operations, 'ttmetal'));
}
getOrLoadDialect(name) {
return this._dialects.get(name);
}
getLoadedDialect(name) {
return this._dialects.has(name);
}
undocumented(name) {
switch (name) {
case 'bstnnx':
case 'common':
case 'dxgml':
case 'dxgml_pattern':
case 'gim':
case 'nir':
case 'nn':
case 'torq_hl':
case 'xir':
case 'xth':
return true;
default:
return false;
}
}
};
_.Dialect = class {
constructor(operations, name) {
this._name = name;
this._operations = new Map();
this._customDirectives = new Map();
this._customTypes = new Map();
this._customAttributes = new Map();
this.registerCustomDirective('DynamicIndexList', this.parseDynamicIndexList.bind(this));
this.registerCustomDirective('NestedDynamicIndexList', this.parseNestedDynamicIndexList.bind(this));
this.registerCustomDirective('Offsets', this.parseOffsets.bind(this));
this.registerCustomDirective('SymbolVisibility', this.parseSymbolVisibility.bind(this));
this.registerCustomDirective('TypeOrAttr', this.parseTypeOrAttr.bind(this));
this.registerCustomDirective('CopyOpRegion', this.parseCopyOpRegion.bind(this));
this.registerCustomDirective('SizeAwareType', this.parseSizeAwareType.bind(this));
this.registerCustomDirective('ResultTypeList', this.parseResultTypeList.bind(this));
this.registerCustomDirective('CmdParameterGatherOperations', this.parseCmdParameterGatherOperations.bind(this));
this.registerCustomDirective('AsyncParameterGatherOperations', this.parseAsyncParameterGatherOperations.bind(this));
this.registerCustomDirective('CmdParameterScatterOperations', this.parseCmdParameterScatterOperations.bind(this));
this.registerCustomDirective('AsyncParameterScatterOperations', this.parseAsyncParameterScatterOperations.bind(this));
this.registerCustomDirective('PipelineAttr', this.parsePipelineAttr.bind(this));
this.registerCustomDirective('KnobsDictionary', this.parseKnobsDictionary.bind(this));
this.registerCustomAttribute('TypedAttrInterface', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('VM_ConstantIntegerValueAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('Util_AnySerializableAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('DenseElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('I32ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('I64ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('F64ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('IndexElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('AnyI32ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('AnyIntElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('StringElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('RankedF32ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('RankedF64ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('RankedI32ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('RankedI64ElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('SignlessIntElementsAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('AnyAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('UnitAttr', this.parseUnitAttr.bind(this));
this.registerCustomAttribute('UnitProp', this.parseUnitAttr.bind(this));
this.registerCustomAttribute('SymbolNameAttr', this.parseSymbolNameAttr.bind(this));
this.registerCustomAttribute('SymbolRefAttr', this.parseSymbolRefAttr.bind(this));
this.registerCustomAttribute('FlatSymbolRefAttr', this.parseFlatSymbolRefAttr.bind(this));
this.registerCustomAttribute('OptionalAttr', this.parseOptionalAttr.bind(this));
this.registerCustomAttribute('OptionalProp', this.parseOptionalAttr.bind(this));
this.registerCustomAttribute('DefaultValuedOptionalAttr', this.parseDefaultValuedOptionalAttr.bind(this));
this.registerCustomAttribute('DefaultValuedAttr', this.parseDefaultValuedAttr.bind(this));
this.registerCustomAttribute('DefaultValuedEnumAttr', this.parseDefaultValuedAttr.bind(this));
this.registerCustomAttribute('DefaultValuedProp', this.parseDefaultValuedAttr.bind(this));
this.registerCustomAttribute('ConfinedAttr', this.parseConfinedAttr.bind(this));
this.registerCustomAttribute('IntValidAlignment', this.parseConfinedAttr.bind(this));
this.registerCustomAttribute('TypeAttrOf', this.parseTypeAttrOf.bind(this));
this.registerCustomAttribute('AnyAttrOf', this.parseAnyAttrOf.bind(this));
this.registerCustomAttribute('ArrayAttr', this.parseArrayAttr.bind(this));
this.registerCustomAttribute('TypedArrayAttrBase', this.parseArrayAttr.bind(this));
this.registerCustomAttribute('DenseI64ArrayAttr', _.DenseI64ArrayAttr.parse);
this.registerCustomAttribute('I64Attr', (parser) => parser.parseAttribute(new _.IntegerType('i64')));
this.registerCustomAttribute('I32Attr', (parser) => parser.parseAttribute(new _.IntegerType('i32')));
this.registerCustomAttribute('I16Attr', (parser) => parser.parseAttribute(new _.IntegerType('i16')));
this.registerCustomAttribute('I8Attr', (parser) => parser.parseAttribute(new _.IntegerType('i8')));
this.registerCustomAttribute('I1Attr', (parser) => parser.parseAttribute(new _.IntegerType('i1')));
this.registerCustomAttribute('SI64Attr', (parser) => parser.parseAttribute(new _.IntegerType('si64')));
this.registerCustomAttribute('SI32Attr', (parser) => parser.parseAttribute(new _.IntegerType('si32')));
this.registerCustomAttribute('UI64Attr', (parser) => parser.parseAttribute(new _.IntegerType('ui64')));
this.registerCustomAttribute('UI32Attr', (parser) => parser.parseAttribute(new _.IntegerType('ui32')));
this.registerCustomAttribute('IndexAttr', (parser) => parser.parseAttribute(new _.IntegerType('index')));
this.registerCustomAttribute('I64Prop', (parser) => parser.parseAttribute(new _.IntegerType('i64')));
this.registerCustomAttribute('I32Prop', (parser) => parser.parseAttribute(new _.IntegerType('i32')));
this.registerCustomAttribute('AlignmentProp', (parser) => parser.parseAttribute(new _.IntegerType('i64')));
this.registerCustomAttribute('F64Attr', (parser) => parser.parseAttribute(new _.FloatType('f64')));
this.registerCustomAttribute('F32Attr', (parser) => parser.parseAttribute(new _.FloatType('f32')));
this.registerCustomAttribute('F16Attr', (parser) => parser.parseAttribute(new _.FloatType('f16')));
this.registerCustomAttribute('BF16Attr', (parser) => parser.parseAttribute(new _.FloatType('bf16')));
this.registerCustomAttribute('StrAttr', (parser) => parser.parseAttribute(new _.PrimitiveType('string')));
this.registerCustomAttribute('TypedStrAttr', this.parseTypedAttrInterface.bind(this));
this.registerCustomAttribute('LevelAttr', (parser) => parser.parseAttribute(new _.IntegerType('index')));
this.registerCustomType('Optional', this.parseOptional.bind(this));
this.registerCustomType('AnyTypeOf', (parser) => parser.parseType());
for (const metadata of operations.get(name) || []) {
this.registerOperandName(metadata.name, metadata);
}
}
get name() {
return this._name;
}
readType(/* reader */) {
throw new mlir.Error(`Bytecode dialect '${this._name}' does not support bytecode types.`);
}
readAttribute(/* reader */) {
throw new mlir.Error(`Dialect '${this._name}' does not support bytecode attributes.`);
}
parseResource(/* entry */) {
throw new mlir.Error(`Dialect '${this._name}' does not support resources.`);
}
getOperation(opName) {
const op = this._operations.get(opName);
if (op && !op.metadata._) {
if (Array.isArray(op.metadata.operands)) {
for (const input of op.metadata.operands) {
if (input && input.type) {
input.type = _.Constraint.parse(input.type);
}
}
}
if (Array.isArray(op.metadata.results)) {
for (const output of op.metadata.results) {
if (output && output.type) {
output.type = _.Constraint.parse(output.type);
}
}
}
if (Array.isArray(op.metadata.attributes)) {
for (const attribute of op.metadata.attributes) {
if (attribute && attribute.type) {
attribute.type = _.Constraint.parse(attribute.type);
}
}
}
if (Array.isArray(op.metadata.regions)) {
for (const region of op.metadata.regions) {
if (region && region.type) {
region.type = _.Constraint.parse(region.type);
}
}
}
if (Array.isArray(op.metadata.traits)) {
for (const trait of op.metadata.traits) {
if (trait && trait.type) {
trait.type = _.Constraint.parse(trait.type);
}
}
}
op.metadata._ = true;
}
return op || null;
}
registerOperandName(opName, metadata) {
const opInfo = new _.RegisteredOperationName(this, opName, metadata);
this._operations.set(opName, opInfo);
}
registerCustomDirective(name, parserFn) {
this._customDirectives.set(name, parserFn);
}
registerCustomType(name, parserFn) {
this._customTypes.set(name, parserFn);
}
registerCustomAttribute(name, parserFn) {
this._customAttributes.set(name, parserFn);
}
createBuildableType(constraint) {
switch (constraint.name) {
case 'AMDGPU_TDMDescriptorType': return new _.Type('!amdgpu.tdm.descriptor');
case 'AnyIRModule': return new _.Type('!transform.any_op');
case 'AnySignlessIntegerOrIndex': return new _.IndexType();
case 'Async_CoroHandleType': return new _.Type('!async.coro.handle');
case 'Async_CoroIdType': return new _.Type('!async.coro.id');
case 'Async_CoroStateType': return new _.Type('!async.coro.state');
case 'Async_GroupType': return new _.async.GroupType();
case 'Async_TokenType': return new _.async.TokenType();
case 'Async_ValueType': return new _.async.ValueType(new _.MemRefType([-1n], new _.FloatType('f32')));
case 'BF16': return new _.FloatType('bf16');
case 'BoolType': return new _.Type('!smt.bool');
case 'CanonicalLoopInfoType': return new _.Type('!omp.canonical_loop_info');
case 'ControlType': return new _.Type('!tf_type.control');
case 'CoreRT_OpHandlerType': return new _.Type('!corert.ophandler');
case 'CoreRT_StringType': return new _.Type('!corert.string');
case 'CoreRT_TensorHandleType': return new _.Type('!corert.tensorhandle');
case 'CUDA_BlasGemmAlgorithm': return new _.Type('!cuda.blas.gemm_algorithm');
case 'CUDA_BlasHandle': return new _.Type('!cuda.blas.handle');
case 'CUDA_Device': return new _.Type('!cuda.device');
case 'CUDA_Event': return new _.Type('!cuda.event');
case 'CUDA_Function': return new _.Type('!cuda.function');
case 'CUDA_Module': return new _.Type('!cuda.module');
case 'CUDA_Stream': return new _.Type('!cuda.stream');
case 'EmitC_PointerType': return new _.Type('!emitc.ptr<i32>');
case 'Executor_HostPtr': return new _.Type('!executor.ptr<host>');
case 'Executor_StrLiteral': return new _.Type('!executor.str_literal');
case 'Executor_Table': return new _.Type('!executor.table<>');
case 'F128': return new _.FloatType('f128');
case 'F16': return new _.FloatType('f16');
case 'F32': return new _.FloatType('f32');
case 'F4E2M1FN': return new _.FloatType('f4E2M1FN');
case 'F64': return new _.FloatType('f64');
case 'F6E2M3FN': return new _.FloatType('f6E2M3FN');
case 'F6E3M2FN': return new _.FloatType('f6E3M2FN');
case 'F80': return new _.FloatType('f80');
case 'F8E3M4': return new _.FloatType('f8E3M4');
case 'F8E4M3': return new _.FloatType('f8E4M3');
case 'F8E4M3B11FNUZ': return new _.FloatType('f8E4M3B11FNUZ');
case 'F8E4M3FN': return new _.FloatType('f8E4M3FN');
case 'F8E4M3FNUZ': return new _.FloatType('f8E4M3FNUZ');
case 'F8E5M2': return new _.FloatType('f8E5M2');
case 'F8E5M2FNUZ': return new _.FloatType('f8E5M2FNUZ');
case 'F8E8M0FNU': return new _.FloatType('f8E8M0FNU');
case 'FLOW_Channel': return new _.Type('!flow.channel');
case 'GPU_AsyncToken': return new _.Type('!gpu.async.token');
case 'GPU_SparseDnTensorHandle': return new _.Type('!gpu.sparse.dntensor_handle');
case 'GPU_SparseSpGEMMOpHandle': return new _.Type('!gpu.sparse.spgemmop_handle');
case 'GPU_SparseSpMatHandle': return new _.Type('!gpu.sparse.spmat_handle');
case 'HostBufferType': return new _.Type('!tfrt_ht.host_buffer');
case 'I1': return new _.IntegerType('i1');
case 'I128': return new _.IntegerType('i128');
case 'I16': return new _.IntegerType('i16');
case 'I32': return new _.IntegerType('i32');
case 'I64': return new _.IntegerType('i64');
case 'I8': return new _.IntegerType('i8');
case 'Ifrt_ArrayType': return new _.Type('!ifrt.array');
case 'Ifrt_ControlType': return new _.Type('!ifrt.control');
case 'Index': return new _.IndexType();
case 'IntType': return new _.Type('!smt.int');
case 'IRDL_AttributeType': return new _.Type('!irdl.attribute');
case 'IRDL_RegionType': return new _.Type('!irdl.region');
case 'IREE_Input_GlobalRefAttr': return new _.Type('!iree_input.global.ref');
case 'LLVM_DefaultPointer': return new _.llvm.LLVMPointerType(0);
case 'LLVM_MetadataType': return new _.Type('!llvm.metadata');
case 'LLVM_PointerGeneric': return new _.llvm.LLVMPointerType(0);
case 'LLVM_PointerGlobal': return new _.llvm.LLVMPointerType(1);
case 'LLVM_PointerShared': return new _.llvm.LLVMPointerType(3);
case 'LLVM_PointerSharedCluster': return new _.llvm.LLVMPointerType(7);
case 'LLVM_PointerTensor': return new _.llvm.LLVMPointerType(6);
case 'LLVM_PointerTensorMemory': return new _.llvm.LLVMPointerType(6);
case 'LLVM_TokenType': return new _.Type('!llvm.token');
case 'MLProgram_TokenType': return new _.Type('!ml_program.token');
case 'MlrtAsyncHandleType': return new _.Type('!mlrt.async_handle');
case 'MlrtFutureType': return new _.Type('!mlrt.future');
case 'MlrtPromiseType': return new _.Type('!mlrt.promise');
case 'MPI_Comm': return new _.Type('!mpi.comm');
case 'MPI_Request': return new _.Type('!mpi.request');
case 'MPI_Retval': return new _.Type('!mpi.retval');
case 'Noisy_I32': return new _.Type('!noisy.i32');
case 'NoneType': return new _.NoneType();
case 'NullPointer': return new _.Type('!iree_codegen.null_pointer');
case 'NVGPU_DeviceAsyncToken': return new _.Type('!nvgpu.device.async.token');
case 'NVGPU_MmaSparseSyncMetadataType': return new _.VectorType([2n], new _.IntegerType('i16'));
case 'OMP_MapBoundsType': return new _.Type('!omp.map_bounds_ty');
case 'OpaqueTensorType': return new _.Type('!tf_type.tensor');
case 'OpenACC_DataBoundsType': return new _.Type('!acc.data_bounds');
case 'OpenACC_DeclareTokenType': return new _.Type('!acc.declare_token');
case 'OpenACC_ParWidthType': return new _.Type('!acc.par_width');
case 'OpenMP_MapBoundsType': return new _.Type('!omp.map_bounds_ty');
case 'PDL_Attribute': return new _.pdl.AttributeType();
case 'PDL_Operation': return new _.pdl.OperationType();
case 'PDL_OperationType': return new _.pdl.OperationType();
case 'PDL_RangeOf': return new _.pdl.RangeType(this.createBuildableType(constraint.args[0]));
case 'PDL_RangeType': return new _.Type('!pdl.range<type>');
case 'PDL_Type': return new _.pdl.TypeType();
case 'PDL_TypeType': return new _.pdl.TypeType();
case 'PDL_Value': return new _.pdl.ValueType();
case 'PDL_ValueType': return new _.pdl.ValueType();
case 'Ptr_PtrType': return new _.ptr.PtrType(null);
case 'ROCDL_V16BF16Type': return new _.VectorType([16n], new _.FloatType('bf16'));
case 'ROCDL_V16F16Type': return new _.VectorType([16n], new _.FloatType('f16'));
case 'ROCDL_V16F32Type': return new _.VectorType([16n], new _.FloatType('f32'));
case 'ROCDL_V2BF16Type': return new _.VectorType([2n], new _.FloatType('bf16'));
case 'ROCDL_V2F16Type': return new _.VectorType([2n], new _.FloatType('f16'));
case 'ROCDL_V2F32Type': return new _.VectorType([2n], new _.FloatType('f32'));
case 'ROCDL_V2I16Type': return new _.VectorType([2n], new _.IntegerType('i16'));
case 'ROCDL_V2I32Type': return new _.VectorType([2n], new _.IntegerType('i32'));
case 'ROCDL_V32BF16Type': return new _.VectorType([32n], new _.FloatType('bf16'));
case 'ROCDL_V32F16Type': return new _.VectorType([32n], new _.FloatType('f16'));
case 'ROCDL_V32F32Type': return new _.VectorType([32n], new _.FloatType('f32'));
case 'ROCDL_V3I32Type': return new _.VectorType([3n], new _.IntegerType('i32'));
case 'ROCDL_V6I32Type': return new _.VectorType([6n], new _.IntegerType('i32'));
case 'ROCDL_V8BF16Type': return new _.VectorType([8n], new _.FloatType('bf16'));
case 'ROCDL_V8F16Type': return new _.VectorType([8n], new _.FloatType('f16'));
case 'ROCDL_V8F32Type': return new _.VectorType([8n], new _.FloatType('f32'));
case 'Shape_ExtentTensorType': return new _.RankedTensorType([_.ShapedType.kDynamic], new _.IndexType());
case 'Shape_ShapeType': return new _.Type('!shape.shape');
case 'Shape_SizeType': return new _.Type('!shape.size');
case 'Shape_ValueShapeType': return new _.Type('!shape.value_shape');
case 'Shape_WitnessType': return new _.Type('!shape.witness');
case 'StaticShapeTensorOf': return new _.UnrankedTensorType(this.createBuildableType(constraint.args[0][0]));
case 'SI16': return new _.IntegerType('si16');
case 'SI32': return new _.IntegerType('si32');
case 'SI64': return new _.IntegerType('si64');
case 'SI8': return new _.IntegerType('si8');
case 'SMTFuncType': return new _.Type('!smt.func<() -> ()>');
case 'SPIRV_AnyPtr': return new _.spirv.PointerType(new _.IntegerType('i32'), 'StorageBuffer');
case 'SPIRV_BFloat16KHR': return new _.FloatType('bf16');
case 'SPIRV_Bool': return new _.IntegerType('i1');
case 'SPIRV_Float16': return new _.FloatType('f16');
case 'SPIRV_Float32': return new _.FloatType('f32');
case 'SPIRV_Int8': return new _.IntegerType('i8');
case 'SPIRV_Int16': return new _.IntegerType('i16');
case 'SPIRV_Int32': return new _.IntegerType('i32');
case 'SPIRV_Int32Vec4': return new _.VectorType([4n], new _.IntegerType('i32'));
case 'Stream_Channel': return new _.Type('!stream.channel');
case 'Stream_Dim': return new _.IndexType();
case 'Stream_Offset': return new _.IndexType();
case 'Stream_Size': return new _.IndexType();
case 'TensorRTRuntime_Context': return new _.Type('!trtrt.context');
case 'TensorRTRuntime_Engine': return new _.Type('!trtrt.engine');
case 'TensorType': return new _.Type('!tfrt_tensor.tensor');
case 'Token': return new _.TokenType();
case 'TF_MLRT_FutureType': return new _.Type('!tf_mlrt.tensor');
case 'TF32': return new _.FloatType('tf32');
case 'TFAllocatorType': return new _.Type('!tfrt_fallback.tf_allocator');
case 'TFDeviceType': return new _.Type('!tf_mlrt.device');
case 'TfeControlType': return new _.Type('!tfe.control');
case 'TfeTokenType': return new _.Type('!tfe.token');
case 'TFFramework_JITCallableType': return new _.Type('!tf_framework.jit_callable');
case 'TFFramework_OpKernelContextType': return new _.Type('!tf_framework.op_kernel_context');
case 'TFL_Control': return new _.Type('!tfl.control');
case 'TFL_Quint8': return new _.Type('!quant.uniform<u8:f32, 1.0>');
case 'TFL_Str': return new _.Type('!tf.string');
case 'TFR_AttrType': return new _.Type('!tfr.attr');
case 'TFR_TensorListType': return new _.Type('!tfr.tensor_list');
case 'TFR_TensorType': return new _.Type('!tfr.tensor');
case 'TFRT_ChainType': return new _.Type('!tfrt.chain');
case 'TFRT_DeviceType': return new _.Type('!tfrt.device');
case 'TFRT_Fallback_TFTensorType': return new _.Type('!tfrt_fallback.tf_tensor');
case 'TFRT_StringType': return new _.Type('!tfrt.string');
case 'TFRT_TensorTypeType': return new _.Type('!tfrt.tensor_type');
case 'TFTensorType': return new _.Type('!tfrt_fallback.tf_tensor');
case 'Torch_BoolType': return new _.Type('!torch.bool');
case 'Torch_DeviceType': return new _.Type('!torch.Device');
case 'Torch_FloatType': return new _.Type('!torch.float');
case 'Torch_GeneratorType': return new _.Type('!torch.Generator');
case 'Torch_IntType': return new _.Type('!torch.int');
case 'Torch_LinearParamsType': return new _.Type('!torch.LinearParams');
case 'Torch_NoneType': return new _.Type('!torch.none');
case 'Torch_NumberType': return new _.Type('!torch.number');
case 'Torch_StringType': return new _.Type('!torch.str');
case 'Transform_AffineMapParamType': return new _.Type('!transform.affine_map');
case 'Transform_AnyOpType': return new _.Type('!transform.any_op');
case 'Transform_AnyParamType': return new _.Type('!transform.any_param');
case 'Transform_AnyValue': return new _.Type('!transform.any_value');
case 'Transform_AnyValueType': return new _.Type('!transform.any_value');
case 'Transform_ParamType': return new _.Type('!transform.param<i64>');
case 'Transform_TypeParamType': return new _.Type('!transform.type');
case 'TransformHandleTypeInterface': return new _.Type('!transform.any_op');
case 'TestTransformTestDialectParamType': return new _.Type('!transform.test_dialect_param');
case 'TS_FixedRankShape': return new _.OpaqueType('ts', `.fixed_rank_shape.${constraint.args[0].name}`);
case 'TS_PartialShape': return new _.Type('!ts.partial_shape');
case 'TS_Shape': return new _.Type('!ts.shape');
case 'TTG_AsyncToken': return new _.Type('!ttg.async.token');
case 'UI16': return new _.IntegerType('ui16');
case 'UI32': return new _.IntegerType('ui32');
case 'UI64': return new _.IntegerType('ui64');
case 'UI8': return new _.IntegerType('ui8');
case 'Util_BufferType': return new _.Type('!util.buffer');
case 'Util_Offset': return new _.IndexType();
case 'Util_Size': return new _.IndexType();
case 'VM_CondValue': return new _.IntegerType('i32');
case 'VM_RefType': return new _.Type('!vm.ref<?>');
case 'XeGPU_Nbarrier': return new _.Type('!xegpu.nbarrier');
case 'XLA_BufferType': return new _.Type('!xla_framework.buffer');
case 'XLAFramework_BufferType': return new _.Type('!xla_framework.buffer');
default: return null;
}
}
inferResultTypes(op, vars) {
const metadata = op.name.getRegisteredInfo()?.metadata;
if (!metadata?.results) {
return;
}
// If type(results) was used in the format, all result types were explicitly parsed
if (vars.has('results') && vars.get('results').types.length > 0) {
if (op.types.length === 0) {
op.addTypes(vars.get('results').types);
}
return;
}
const operandNames = metadata.operands?.map((o) => o.name) || [];
// Build result types in metadata order
const orderedTypes = [];
// Track current offset in op.types for custom directive handling
let customDirectiveOffset = 0;
// Iterate over each result in metadata order
for (const resultInfo of metadata.results) {
// If this result was parsed from assembly format, use those types
if (vars.has(resultInfo.name) && vars.get(resultInfo.name).types.length > 0) {
const varTypes = vars.get(resultInfo.name).types;
orderedTypes.push(...varTypes);
customDirectiveOffset += varTypes.length;
continue;
}
const mnemonic = typeof resultInfo.type === 'string' ? resultInfo.type : resultInfo.type?.name;
const isVariadicOrOptional = mnemonic === 'Variadic' || mnemonic === 'Optional';
// For Variadic/Optional not in vars, check if custom directive added types to op.types
if (isVariadicOrOptional && op.types.length > customDirectiveOffset) {
// Count how many types belong to this variadic result
// (all types from customDirectiveOffset to end, minus types for later results in vars)
let typesForLaterResults = 0;
for (let i = metadata.results.indexOf(resultInfo) + 1; i < metadata.results.length; i++) {
const laterResult = metadata.results[i];
if (vars.has(laterResult.name) && vars.get(laterResult.name).types.length > 0) {
typesForLaterResults += vars.get(laterResult.name).types.length;
}
}
const variadicCount = op.types.length - customDirectiveOffset - typesForLaterResults;
if (variadicCount > 0) {
for (let i = 0; i < variadicCount; i++) {
orderedTypes.push(op.types[customDirectiveOffset + i]);
}
customDirectiveOffset += variadicCount;
continue;
}
}
// Try to resolve from traits
let resolved = false;
if (metadata.traits) {
for (const trait of metadata.traits) {
if (resolved) {
break;
}
// Handle both string and object trait types
const traitName = trait.type?.name;
const traitArgs = trait.type?.args;
// SameOperandsAndResultType: result type = first operand type
if (traitName === 'SameOperandsAndResultType') {
if (op.operands.length > 0 && op.operands[0].type) {
orderedTypes.push(op.operands[0].type);
resolved = true;
}
}
// AllTypesMatch: result type = matched operand/result/attribute type
if (traitName === 'AllTypesMatch') {
const names = traitArgs?.[0];
if (Array.isArray(names) && (names.includes(resultInfo.name) || (resultInfo.name === 'result' && names.includes('result')))) {
const sourceTypes = [];
for (const argName of names) {
if (argName === resultInfo.name) {
continue;
}
// Check if source is in vars (another parsed result)
if (vars.get(argName) && vars.get(argName).types.length > 0) {
sourceTypes.push(...vars.get(argName).types);
break;
}
// Check if source is an operand
const operandMetaIdx = operandNames.indexOf(argName);
if (operandMetaIdx >= 0) {
const operandMeta = metadata.operands[operandMetaIdx];
const operandTypeObj = operandMeta?.type;
const operandTypeName = typeof operandTypeObj === 'string' ? operandTypeObj : (operandTypeObj?.name || '');
const isVariadic = operandTypeName === 'Variadic' || operandTypeName.startsWith('Variadic<');
if (isVariadic) {
for (let j = 0; j < op.operands.length; j++) {
if (op.operands[j].type) {
sourceTypes.push(op.operands[j].type);
}
}
} else if (operandMetaIdx < op.operands.length && op.operands[operandMetaIdx].type) {
sourceTypes.push(op.operands[operandMetaIdx].type);
}
if (sourceTypes.length > 0) {
break;
}
}
// Check if source is an attribute with a type
if (metadata.attributes) {
const attrMeta = metadata.attributes.find((a) => a.name === argName);
if (attrMeta) {
const attr = op.attributes.get(argName);
if (attr && attr.type) {
sourceTypes.push(attr.type);
break;
}
}
}
}
if (sourceTypes.length > 0) {
orderedTypes.push(...sourceTypes);
resolved = true;
}
}
}
// TypesMatchWith: result type = transform(source type)
if (traitName === 'TypesMatchWith') {
const [lhsArg, rhsArg, transformer] = traitArgs || [];
let resultArg = null;
let sourceArg = null;
if ((rhsArg === resultInfo.name || (resultInfo.name === 'result' && rhsArg === 'result'))) {
resultArg = rhsArg;
sourceArg = lhsArg;
} else if ((lhsArg === resultInfo.name || (resultInfo.name === 'result' && lhsArg === 'result'))) {
resultArg = lhsArg;
sourceArg = rhsArg;
}
if (resultArg && sourceArg) {
const sourceTypes = [];
// Check if source is in vars (parsed result or operand type from type($var))
if (vars.has(sourceArg) && vars.get(sourceArg).types.length > 0) {
sourceTypes.push(...vars.get(sourceArg).types);
}
// Check if source is an operand - fall back to operand's value type
if (sourceTypes.length === 0) {
const operands = metadata.operands || [];
let actualOperandIdx = 0;
for (let i = 0; i < operands.length; i++) {
const operandTypeObj = operands[i].type;
const operandType = typeof operandTypeObj === 'string' ? operandTypeObj : (operandTypeObj?.name || '');
const isEnum = /\{[^}]+\}/.test(operandType) || operandType.startsWith('AtomicBinOp') || operandType.startsWith('AtomicOrdering') || operandType.startsWith('LLVM_AtomicOrdering') || operandType.includes('Enum');
if (operands[i].name === sourceArg) {
const isVariadicSource = operandType === 'Variadic';
if (isVariadicSource) {
for (let j = actualOperandIdx; j < op.operands.length; j++) {
if (op.operands[j].type) {
sourceTypes.push(op.operands[j].type);
}
}
} else if (actualOperandIdx < op.operands.length && op.operands[actualOperandIdx].type) {
sourceTypes.push(op.operands[actualOperandIdx].type);
}
break;
}
if (!isEnum) {
actualOperandIdx++;
}
}
}
if (sourceTypes.length > 0) {
for (const sourceType of sourceTypes) {
let resultType = sourceType;
if (transformer === '::getI1SameShape($_self)') {
if (sourceType instanceof _.VectorType) {
resultType = new _.VectorType(sourceType.dimensions, new _.IntegerType('i1'), sourceType.scalableDims);
} else if (sourceType instanceof _.RankedTensorType) {
resultType = new _.RankedTensorType(sourceType.shape, new _.IntegerType('i1'));
} else {
resultType = new _.IntegerType('i1');
}
} else if (transformer && transformer.includes('.getElementType()')) {
if (sourceType && sourceType.elementType) {
resultType = sourceType.elementType;
} else {
throw new mlir.Error(`Cannot get element type.`);
}
} else if (transformer && transformer.includes('getPointeeType')) {
if (sourceType && sourceType.pointeeType) {
resultType = sourceType.pointeeType;
} else {
throw new mlir.Error(`Cannot get pointee type.`);
}
} else if (transformer && transformer.includes('getValAndBoolStructType')) {
const typeStr = sourceType.toString ? sourceType.toString() : String(sourceType);
resultType = new _.Type(`!llvm.struct<(${typeStr}, i1)>`);
} else if (transformer && transformer.includes('.getResults()')) {
// Extract results from a FunctionType
if (sourceType instanceof _.FunctionType && sourceType.results) {
orderedTypes.push(...sourceType.results);
resolved = true;
continue; // Already added results, skip single push below
}
}
if (resultType) {
orderedTypes.push(resultType);
resolved = true;
}
}
}
}
}
// FirstAttrDerivedResultType: result type = first attribute's type
if (traitName === 'FirstAttrDerivedResultType') {
const firstAttr = metadata.attributes?.[0];
if (firstAttr) {
const attr = op.attributes.get(firstAttr.name);
if (attr && attr.type) {
orderedTypes.push(attr.type);
resolved = true;
}
}
}
}
}
// For single variadic result matching single variadic operand with same element type, infer from operands
if (!resolved && resultInfo.type?.name === 'Variadic' &&
metadata.results?.length === 1 && metadata.operands?.length === 1) {
const resultElementType = resultInfo.type.args?.[0]?.name;
const operandInfo = metadata.operands[0];
if (operandInfo.type?.name === 'Variadic' &&
operandInfo.type.args?.[0]?.name === resultElementType) {
for (const operand of op.operands) {
if (operand.type) {
orderedTypes.push(operand.type);
}
}
resolved = orderedTypes.length > 0;
}
}
// Fallback: try buildable type (skip for variadic/optional)
if (!resolved && resultInfo.type && !isVariadicOrOptional) {
const type = this.createBuildableType(resultInfo.type);
if (type) {
orderedTypes.push(type);
}
}
}
// Replace op.types with correctly ordered types
if (orderedTypes.length > 0) {
op.types.length = 0;
op.addTypes(orderedTypes);
}
}
parseDirective(directive, parser, op, opInfo, directives, i, vars) {
const isVariadic = (type) => {
if (type.name === 'Variadic' || type.name === 'VariadicOfVariadic') {
return true;
}
if (Array.isArray(type.args) && type.args.length > 0) {
return isVariadic(type.args[0]);
}
return false;
};
const isVariadicOfVariadic = (type) => {
if (type.name === 'VariadicOfVariadic') {
return true;
}
if (Array.isArray(type.args) && type.args.length > 0) {
return isVariadicOfVariadic(type.args[0]);
}
return false;
};
const isOptional = (type) => {
if (type.name === 'Optional') {
return true;
}
if (Array.isArray(type.args) && type.args.length > 0) {
return isOptional(type.args[0]);
}
return false;
};
const peek = (...kinds) => {
const token = parser.parser.getToken();
return kinds.some((kind) => kind === 'keyword' ? token.isKeyword() : token.is(kind));
};
switch (directive.type) {
case 'whitespace':
// Skip whitespace directives - they're just formatting hints
break;
case 'literal':
// Make '[' optional when followed by a variadic operand with buildable type
if (directive.value === '[') {
if (!parser.parseOptionalLSquare()) {
const nextDir = directives[i + 1];
if (nextDir && nextDir.type === 'operand_ref') {
const operandMeta = opInfo.metadata?.operands?.find((o) => o.name === nextDir.name);
if (operandMeta && isVariadic(operandMeta.type)) {
vars.set('_skipClosingBracket', true);
break;
}
}
parser.parseLSquare();
}
break;
}
if (directive.value === ']' && vars.get('_skipClosingBracket')) {
vars.delete('_skipClosingBracket');
break;
}
switch (directive.value) {
case '(': parser.parseLParen(); break;
case ')': parser.parseRParen(); break;
case '[': parser.parseLSquare(); break;
case ']': parser.parseRSquare(); break;
case '{': parser.parseLBrace(); break;
case '}': parser.parseRBrace(); break;
case ':': parser.parseColon(); break;
case '=': parser.parseEqual(); break;
case ',': parser.parseComma(); break;
case '<': parser.parseLess(); break;
case '>': parser.parseGreater(); break;
case '->': parser.parseArrow(); break;
case '*': parser.parseStar(); break;
case '+': parser.parsePlus(); break;
case '?': parser.parseQuestion(); break;
case '|': parser.parseVerticalBar(); break;
case '...': parser.parseEllipsis(); break;
default: parser.parseKeyword(directive.value); break;
}
break;
case 'region_ref': {
const regionMeta = opInfo.metadata && opInfo.metadata.regions && opInfo.metadata.regions.find((r) => r.name === directive.name);
const isVariadicRegion = regionMeta && regionMeta.type && regionMeta.type.name === 'VariadicRegion';
const isIsolated = op.name.getRegisteredInfo().hasTrait('IsolatedFromAbove');
if (isVariadicRegion) {
if (parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
while (parser.parseOptionalComma()) {
if (!parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
break;
}
}
}
} else {
const region = op.addRegion();
parser.parseRegion(region, undefined, isIsolated);
}
break;
}
case 'successor_ref': {
if (!op.successors) {
op.successors = [];
}
// Check if this successor is variadic from metadata or context
const refName = directive.name;
let isVariadicSuccessor = false;
if (opInfo.metadata && opInfo.metadata.successors) {
const successorMeta = opInfo.metadata.successors.find((s) => s.name === refName);
if (successorMeta && successorMeta.type) {
// Check for VariadicSuccessor type
const typeStr = typeof successorMeta.type === 'string' ? successorMeta.type : successorMeta.type.name;
isVariadicSuccessor = typeStr && typeStr.startsWith('VariadicSuccessor');
}
}
// Also check context: if next directive is ')' literal, we're inside parentheses
const nextDir = i + 1 < directives.length ? directives[i + 1] : null;
const isVariadicContext = isVariadicSuccessor || (nextDir && nextDir.type === 'literal' && nextDir.value === ')');
// Check if next directive is a custom directive that handles the parentheses (like ResultTypeList)
const nextDirHandlesParens = nextDir && nextDir.type === 'custom' &&
(nextDir.parser === 'ResultTypeList' || nextDir.parser === 'TypeList');
const parseOneSuccessor = () => {
const successor = parser.parseSuccessor();
// Don't consume '(' if next directive handles it (e.g., custom<ResultTypeList>)
if (!nextDirHandlesParens && parser.parseOptionalLParen()) {
successor.arguments = [];
let operand = parser.parseOptionalOperand();
while (operand) {
successor.arguments.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.resolveOperands(successor.arguments, parser.parseOptionalColonTypeList());
parser.parseOptionalRParen();
}
op.successors.push(successor);
};
if (isVariadicContext) {
// Variadic successors: parse 0 or more successors
let successor = parser.parseOptionalSuccessor();
while (successor) {
// Don't consume '(' if next directive handles it
if (!nextDirHandlesParens && parser.parseOptionalLParen()) {
successor.arguments = [];
let operand = parser.parseOptionalOperand();
while (operand) {
successor.arguments.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.resolveOperands(successor.arguments, parser.parseOptionalColonTypeList());
parser.parseOptionalRParen();
}
op.successors.push(successor);
if (!parser.parseOptionalComma()) {
break;
}
successor = parser.parseOptionalSuccessor();
}
} else {
parseOneSuccessor();
}
break;
}
case 'attribute_ref': {
const refName = directive.name;
if (op.attributes.has(refName)) {
break;
}
const attrInfo = opInfo.metadata && opInfo.metadata.attributes && opInfo.metadata.attributes.find((attr) => attr.name === refName);
const attrType = attrInfo ? attrInfo.type : null;
let attrValue = null;
// Pass type to suppress : type suffix parsing (it's a separate directive in assembly format)
if (attrType && attrType !== 'Attribute') {
attrValue = this.parseCustomAttributeWithFallback(parser, attrType);
} else {
attrValue = parser.parseAttribute(attrType || 'Attribute');
}
if (attrValue) {
op.addAttribute(refName, attrValue);
}
break;
}
case 'operand_ref': {
const name = directive.name;
const input = opInfo.metadata?.operands?.find((inp) => inp.name === name);
const isVariadicOp = input ? isVariadic(input.type) : false;
const isVariadicOfVariadicOp = input ? isVariadicOfVariadic(input.type) : false;
const isOptionalOp = input ? isOptional(input.type) : false;
// Check for buildable types using createType
let buildableType = null;
if (isVariadicOp && input?.type?.args?.[0]) {
buildableType = this.createBuildableType(input.type.args[0]);
} else if (input?.type) {
buildableType = this.createBuildableType(input.type);
}
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
const entry = vars.get(name);
if (isVariadicOfVariadicOp) {
// Parse grouped operands: (op, op), (), (op)
do {
if (!parser.parseOptionalLParen()) {
break;
}
let operand = parser.parseOptionalOperand();
while (operand) {
entry.operands.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.parseRParen();
} while (parser.parseOptionalComma());
} else if (isVariadicOp) {
for (;;) {
const operand = parser.parseOptionalOperand();
if (operand) {
entry.operands.push(operand);
if (buildableType) {
entry.types.push(buildableType);
}
if (!parser.parseOptionalComma()) {
break;
}
} else if (buildableType) {
// Handle integer literals for buildable integer types (e.g., I32)
let value = '';
if (parser.parseOptionalMinus()) {
value = '-';
}
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
value += intVal;
entry.operands.push({ name: value, literal: true });
if (buildableType) {
entry.types.push(buildableType);
}
if (!parser.parseOptionalComma()) {
break;
}
} else {
break;
}
}
} else {
const operand = parser.parseOptionalOperand();
if (operand) {
entry.operands.push(operand);
if (buildableType) {
entry.types.push(buildableType);
}
} else if (buildableType) {
// Handle integer literals for buildable integer types (e.g., I32)
let value = '';
if (parser.parseOptionalMinus()) {
value = '-';
}
const intVal = parser.parseOptionalInteger();
if (intVal !== null) { // eslint-disable-line no-negated-condition
value += intVal;
entry.operands.push({ name: value, literal: true });
entry.types.push(buildableType);
} else {
// Check if this is a region, not an operand
const isActualOperand = opInfo.metadata?.operands?.some((inp) => inp.name === name);
if (isActualOperand) {
// Try parsing as symbol name (@identifier)
const symbolName = parser.parseOptionalSymbolName();
if (symbolName) {
op.addAttribute(name, symbolName.value || symbolName);
} else if (!isOptionalOp) {
// Try keyword (enum value)
const inputType = input?.type;
if (inputType && Array.isArray(inputType.values)) {
const keyword = parser.parseKeyword();
op.addAttribute(name, keyword);
} else {
// Try integer
const integerVal = parser.parseOptionalInteger();
if (integerVal !== null) { // eslint-disable-line no-negated-condition
op.addAttribute(name, String(integerVal));
} else {
// Fallback: try parseAttribute
const attr = parser.parseOptionalAttribute();
if (attr) {
op.addAttribute(name, attr);
}
}
}
}
} else {
const regionMeta = opInfo.metadata?.regions?.find((r) => r.name === name);
const isVariadicRegion = regionMeta?.type?.name === 'VariadicRegion';
const isIsolated = op.name.hasTrait('IsolatedFromAbove');
if (isVariadicRegion) {
if (parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
while (parser.parseOptionalComma()) {
if (!parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
break;
}
}
}
} else {
parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated);
}
}
}
} else {
// Check if this is a region, not an operand
const isActualOperand = opInfo.metadata?.operands?.some((inp) => inp.name === name);
if (isActualOperand) {
// Try parsing as symbol name (@identifier)
const symbolName = parser.parseOptionalSymbolName();
if (symbolName) {
op.addAttribute(name, symbolName.value || symbolName);
} else if (!isOptionalOp) {
// Try keyword (enum value)
const inputType = input?.type;
if (inputType && Array.isArray(inputType.values)) {
const keyword = parser.parseKeyword();
op.addAttribute(name, keyword);
} else {
// Try integer
const integerVal = parser.parseOptionalInteger();
if (integerVal !== null) { // eslint-disable-line no-negated-condition
op.addAttribute(name, String(integerVal));
} else {
// Fallback: try parseAttribute
const attr = parser.parseOptionalAttribute();
if (attr) {
op.addAttribute(name, attr);
}
}
}
}
} else {
const regionMeta = opInfo.metadata?.regions?.find((r) => r.name === name);
const isVariadicRegion = regionMeta?.type?.name === 'VariadicRegion';
const isIsolated = op.name.hasTrait('IsolatedFromAbove');
if (isVariadicRegion) {
if (parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
while (parser.parseOptionalComma()) {
if (!parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
break;
}
}
}
} else {
parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated);
}
}
}
}
break;
}
case 'operands': {
if (!vars.has('operands')) {
vars.set('operands', { operands: [], types: [] });
}
const operandsList = parser.parseOperandList();
vars.get('operands').operands.push(...operandsList);
break;
}
case 'results': {
const args = parser.parseArgumentList('none', true);
const types = args.map((a) => a.type).filter((t) => t);
op.addTypes(types);
break;
}
case 'type':
case 'qualified': {
if (!directive.args || directive.args.length === 0) {
// Bare type directive - parse types for operands
const types = parser.parseTypeList();
if (vars.has('operands')) {
vars.get('operands').types.push(...types);
}
break;
}
const arg = directive.args[0] === 'type' && directive.args.length > 1 ? directive.args[1] : directive.args[0];
// Handle qualified($attr) - attribute reference
if (directive.type === 'qualified' && arg.startsWith('$') && !arg.startsWith('$results') && !arg.startsWith('$operands')) {
if (!arg.startsWith('type($')) {
const attrName = arg.substring(1);
const attr = parser.parseAttribute();
if (attr) {
// Store full attribute object to preserve .type for FirstAttrDerivedResultType
op.addAttribute(attrName, attr);
}
break;
}
}
// Extract name from $name or type($name) or type(operands) or type(results)
let name = null;
if (arg.startsWith('type($') && arg.endsWith(')')) {
name = arg.substring(6, arg.length - 1);
} else if (arg.startsWith('type(') && arg.endsWith(')')) {
// Handle type(operands) or type(results)
name = arg.substring(5, arg.length - 1);
} else if (arg.startsWith('$')) {
name = arg.substring(1);
} else if (arg === 'results' || arg === 'operands') {
name = arg;
}
if (!name) {
break;
}
// Check if it's a result or operand
const resultMeta = opInfo.metadata?.results?.find((r) => r.name === name);
const operandMeta = opInfo.metadata?.operands?.find((o) => o.name === name);
const isResult = Boolean(resultMeta) && !operandMeta;
let isVariadicType = false;
if (resultMeta) {
isVariadicType = isVariadic(resultMeta.type);
} else if (operandMeta) {
isVariadicType = isVariadic(operandMeta.type);
}
const isVariadicOfVariadicType = operandMeta ? isVariadicOfVariadic(operandMeta.type) : false;
let isOptionalType = false;
if (operandMeta) {
isOptionalType = isOptional(operandMeta.type);
} else if (resultMeta) {
isOptionalType = isOptional(resultMeta.type);
}
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
const entry = vars.get(name);
if (isResult || name === 'results') {
// Result type - add to op.types and track in entry.types
if (isVariadicType || name === 'results') {
const types = parser.parseTypeList();
entry.types.push(...types);
op.addTypes(types);
} else if (isOptionalType && op.types.length === 0) {
const type = parser.parseOptionalType();
if (type) {
entry.types.push(type);
op.addTypes([type]);
}
} else {
const type = this.parseCustomTypeWithFallback(parser, resultMeta?.type);
entry.types.push(type);
op.addTypes([type]);
}
} else if (isVariadicOfVariadicType) {
// Parse grouped types: (type, type), (), (type)
do {
if (!parser.parseOptionalLParen()) {
break;
}
if (!parser.parseOptionalRParen()) {
entry.types.push(...parser.parseTypeList());
parser.parseRParen();
}
} while (parser.parseOptionalComma());
} else if (isVariadicType || name === 'operands') {
// Variadic operand - parse type list
entry.types.push(...parser.parseTypeList());
} else if (entry.operands.length > 0) {
// Single operand - parse one type per operand
// Clear any previously inferred (buildable) types since explicit type is parsed
const type = this.parseCustomTypeWithFallback(parser, operandMeta?.type);
entry.types = [];
for (let j = 0; j < entry.operands.length; j++) {
entry.types.push(type);
}
} else if (isOptionalType) {
// Optional operand - parse type if present
const type = parser.parseOptionalType();
if (type) {
entry.types.push(type);
}
} else {
// No operands collected yet, just parse and store type
// Use custom type parser if available for the operand/result constraint
const typeConstraint = operandMeta?.type || resultMeta?.type;
const type = this.parseCustomTypeWithFallback(parser, typeConstraint);
entry.types.push(type);
}
break;
}
case 'attr_dict_with_keyword':
parser.parseOptionalAttrDictWithKeyword(op.attributes);
break;
case 'attr_dict':
parser.parseOptionalAttrDict(op.attributes);
break;
case 'prop_dict':
if (parser.parseOptionalLess()) {
op.propertiesAttr = parser.parseAttribute();
parser.parseGreater();
}
break;
case 'regions': {
const isIsolated = op.name.hasTrait('IsolatedFromAbove');
while (parser.parseOptionalRegion(op.addRegion(), undefined, isIsolated)) {
if (!parser.parseOptionalComma()) {
break;
}
}
break;
}
case 'successors': {
op.successors = op.successors || [];
let succ = parser.parseOptionalSuccessor();
if (succ) {
op.successors.push(succ);
while (parser.parseOptionalComma()) {
succ = parser.parseOptionalSuccessor();
if (succ) {
op.successors.push(succ);
} else {
break;
}
}
}
break;
}
case 'functional_type': {
const typeLoc = parser.getCurrentLocation().copy();
const type = parser.parseType();
if (type instanceof _.FunctionType === false) {
parser.emitError(typeLoc, 'Expected function type');
}
// Distribute input types to operands in metadata order
let typeIndex = 0;
for (const operandMeta of opInfo.metadata?.operands || []) {
if (vars.has(operandMeta.name)) {
const entry = vars.get(operandMeta.name);
for (let j = 0; j < entry.operands.length && typeIndex < type.inputs.length; j++) {
entry.types.push(type.inputs[typeIndex]);
typeIndex++;
}
}
}
op.addTypes(type.results);
// Track result types - second arg is typically the result (e.g., $outputs or results)
if (directive.args && directive.args.length > 1) {
const resultArg = directive.args[1];
if (resultArg) {
// Handle both $name and name (without $)
const resultName = resultArg.startsWith('$') ? resultArg.substring(1) : resultArg;
if (resultName === 'results') {
// 'results' is a keyword meaning all results - create vars entry
if (!vars.has('results')) {
vars.set('results', { types: [] });
}
vars.get('results').types.push(...type.results);
} else if (vars.has(resultName)) {
vars.get(resultName).types.push(...type.results);
}
}
}
break;
}
case 'custom': {
const fn = this._customDirectives.get(directive.parser);
if (!fn) {
throw new mlir.Error(`Custom directive parser '${directive.parser}' not implemented.`);
}
// Parse args and pass resolved arrays
// Always pass parser and op first, then resolved args
const callArgs = [parser, op];
for (const arg of directive.args || []) {
if (arg.startsWith('ref(type($') && arg.endsWith('))')) {
// ref(type($name)) - reference to the type of an operand
// 'ref(type($' is 10 characters, so slice from index 10 to skip -2 for '))'
const name = arg.slice(10, -2);
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
callArgs.push(vars.get(name).types);
} else if (arg.startsWith('ref($') && arg.endsWith(')')) {
const name = arg.slice(5, -1);
// Check if this is an attribute reference
const isAttribute = opInfo.metadata?.attributes?.some((a) => a.name === name);
if (isAttribute) {
// Get attribute value from op.attributes
const attrValue = op.attributes.get(name);
callArgs.push(attrValue);
} else {
// Operand reference - get from ctx
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
callArgs.push(vars.get(name).operands);
}
} else if (arg.startsWith('type($') && arg.endsWith(')')) {
const name = arg.slice(6, -1);
// Always use vars entry for the specific name (result or operand)
// This ensures Optional results like asyncToken are tracked separately
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
callArgs.push(vars.get(name).types);
} else if (arg.startsWith('$')) {
const name = arg.slice(1);
// Could be operand ref or attribute name
const isOperand = opInfo.metadata?.operands?.some((o) => o.name === name);
if (isOperand) {
if (!vars.has(name)) {
vars.set(name, { operands: [], types: [] });
}
callArgs.push(vars.get(name).operands);
} else {
// Pass attribute name for custom directive to handle
callArgs.push(name);
}
} else {
// Warn if a $-prefixed arg wasn't resolved - indicates missing metadata
if (typeof arg === 'string' && arg.startsWith('$')) {
throw new mlir.Error(`Custom directive '${directive.parser}' received unresolved arg '${arg}' for op '${opInfo.metadata.name}'. Check metadata for missing operand/attribute definition.`);
}
callArgs.push(arg);
}
}
fn(...callArgs);
break;
}
case 'oilist': {
const clauses = directive.content.split('|').map((c) => c.trim());
const parsedClauses = [];
for (const clauseStr of clauses) {
const clauseParser = new _.AssemblyFormatParser({ ...opInfo.metadata, assemblyFormat: clauseStr });
const elements = clauseParser.parse();
parsedClauses.push({ elements, parsed: false, clauseStr });
}
// Helper to check if a clause's variables are used by later custom directives
const isHandledByCustomDirective = (clauseStr) => {
const varMatches = clauseStr.matchAll(/\$(\w+)/g);
const clauseVars = [...varMatches].map((m) => m[1]);
if (clauseVars.length === 0) {
return false;
}
for (let j = i + 1; j < directives.length; j++) {
const laterDir = directives[j];
if (laterDir.type === 'custom' && laterDir.args && Array.isArray(laterDir.args)) {
const customVarNames = [];
for (const arg of laterDir.args) {
const argVarMatches = arg.matchAll(/\$(\w+)/g);
for (const match of argVarMatches) {
customVarNames.push(match[1]);
}
}
if (clauseVars.some((v) => customVarNames.includes(v))) {
return true;
}
}
}
return false;
};
let progress = true;
while (progress) {
progress = false;
for (const clause of parsedClauses) {
if (clause.parsed) {
continue;
}
if (clause.elements.length === 0) {
continue;
}
if (isHandledByCustomDirective(clause.clauseStr)) {
clause.parsed = true;
continue;
}
const firstElem = clause.elements.find((elem) => elem.type !== 'whitespace');
let matches = false;
if (firstElem && firstElem.type === 'literal') {
switch (firstElem.value) {
case '(': matches = parser.parseOptionalLParen(); break;
case ')': matches = parser.parseOptionalRParen(); break;
case '[': matches = parser.parseOptionalLSquare(); break;
case ']': matches = parser.parseOptionalRSquare(); break;
case '{': matches = parser.parseOptionalLBrace(); break;
case '}': matches = parser.parseOptionalRBrace(); break;
case ':': matches = parser.parseOptionalColon(); break;
case '=': matches = parser.parseOptionalEqual(); break;
case ',': matches = parser.parseOptionalComma(); break;
case '<': matches = parser.parseOptionalLess(); break;
case '>': matches = parser.parseOptionalGreater(); break;
case '->': matches = parser.parseOptionalArrow(); break;
case '...': matches = parser.parseOptionalEllipsis(); break;
default: matches = parser.parseOptionalKeyword(firstElem.value); break;
}
}
if (matches) {
const firstElemIdx = clause.elements.indexOf(firstElem) + 1;
for (let elemIdx = firstElemIdx; elemIdx < clause.elements.length; elemIdx++) {
this.parseDirective(clause.elements[elemIdx], parser, op, opInfo, clause.elements, elemIdx, vars);
}
clause.parsed = true;
progress = true;
}
}
}
break;
}
case 'optional_group': {
let shouldParse = false;
const firstElem = directive.elements.find((elem) => elem.type !== 'whitespace');
if (firstElem) {
if (firstElem.type === 'literal') {
switch (firstElem.value) {
case '(': shouldParse = parser.parseOptionalLParen(); break;
case ')': shouldParse = parser.parseOptionalRParen(); break;
case '[': shouldParse = parser.parseOptionalLSquare(); break;
case ']': shouldParse = parser.parseOptionalRSquare(); break;
case '{': shouldParse = parser.parseOptionalLBrace(); break;
case '}': shouldParse = parser.parseOptionalRBrace(); break;
case ':': shouldParse = parser.parseOptionalColon(); break;
case '=': shouldParse = parser.parseOptionalEqual(); break;
case ',': shouldParse = parser.parseOptionalComma(); break;
case '<': shouldParse = parser.parseOptionalLess(); break;
case '>': shouldParse = parser.parseOptionalGreater(); break;
case '?': shouldParse = parser.parseOptionalQuestion(); break;
case '->': shouldParse = parser.parseOptionalArrow(); break;
case '...': shouldParse = parser.parseOptionalEllipsis(); break;
default: shouldParse = parser.parseOptionalKeyword(firstElem.value); break;
}
if (shouldParse) {
shouldParse = 'skip_first';
}
} else if (firstElem.type === 'attribute_ref') {
const attrInfo = opInfo.metadata && opInfo.metadata.attributes && opInfo.metadata.attributes.find((attr) => attr.name === firstElem.name);
const attrType = attrInfo ? attrInfo.type : null;
// Check if attribute type is an array (TypedArrayAttrBase, ArrayAttr, etc.)
// Handle both string types ("OptionalAttr<StrAttr>") and object types ({name: "OptionalAttr", args: [...]})
let elementType = attrType;
const getTypeName = (t) => typeof t === 'string' ? t : t?.name;
const typeContains = (t, pattern) => {
const name = typeof t === 'string' ? t : t?.name;
return name && name.includes(pattern);
};
// Unwrap wrapper types
if (typeContains(elementType, 'OptionalAttr')) {
elementType = typeof elementType === 'string' ? elementType.replace(/OptionalAttr<(.+)>/, '$1') : elementType.args?.[0];
}
if (typeContains(elementType, 'DefaultValuedAttr') || typeContains(elementType, 'DefaultValuedStrAttr')) {
elementType = typeof elementType === 'string' ? elementType.replace(/DefaultValuedAttr<(.+?)(?:,.*?)?>/, '$1') : elementType.args?.[0];
}
if (typeContains(elementType, 'ConfinedAttr')) {
elementType = typeof elementType === 'string' ? elementType.replace(/ConfinedAttr<(.+?)(?:,.*?)?>/, '$1') : elementType.args?.[0];
}
const mnemonic = getTypeName(elementType);
let shouldTryParse = false;
if (typeContains(elementType, 'TypedArrayAttrBase') || typeContains(elementType, 'ArrayAttr')) {
shouldTryParse = peek(_.Token.l_square);
} else if (mnemonic && /^[SU]?I\d+Attr$|^IntegerAttr$|^IndexAttr$/.test(mnemonic)) {
shouldTryParse = peek(_.Token.integer, _.Token.minus);
} else if (typeContains(elementType, 'ElementsAttr')) {
shouldTryParse = peek(_.Token.kw_dense, _.Token.kw_sparse, _.Token.kw_array, _.Token.kw_dense_resource);
} else if (mnemonic === 'StrAttr' || mnemonic === 'StringAttr') {
shouldTryParse = peek(_.Token.string);
} else {
shouldTryParse = peek(_.Token.bare_identifier, 'keyword', _.Token.hash_identifier, _.Token.at_identifier, _.Token.string, _.Token.l_square, _.Token.integer);
}
if (shouldTryParse) {
let result = null;
if (attrType && attrType !== 'Attribute') {
result = this.parseCustomAttributeWithFallback(parser, attrType);
} else {
result = parser.parseOptionalAttribute(attrType || 'Attribute');
}
if (result !== null) {
op.addAttribute(firstElem.name, result);
shouldParse = true;
}
}
} else if (firstElem.type === 'successor_ref') {
shouldParse = peek(_.Token.caret_identifier);
} else if (firstElem.type === 'region_ref') {
shouldParse = peek(_.Token.l_brace);
} else if (firstElem.type === 'operand_ref') {
let isKeywordInput = false;
if (opInfo.metadata && opInfo.metadata.operands) {
const inputInfo = opInfo.metadata.operands.find((inp) => inp.name === firstElem.name);
if (inputInfo) {
const inputType = inputInfo.type;
if (typeof inputType === 'string' && (inputType.includes('Prop') || inputType.endsWith('Predicate') || inputType.includes('Flags') || inputType.includes('Enum'))) {
isKeywordInput = true;
}
}
}
if (isKeywordInput) {
shouldParse = peek(_.Token.bare_identifier);
} else {
shouldParse = peek(_.Token.percent_identifier);
}
} else if (firstElem.type === 'operands') {
shouldParse = peek(_.Token.l_paren, _.Token.percent_identifier);
} else if (firstElem.type === 'custom') {
const fn = this._customDirectives.get(firstElem.parser);
if (fn) {
const resolvedArgs = (firstElem.args || []).map((arg) => typeof arg === 'string' && arg.startsWith('$') ? arg.slice(1) : arg);
const result = fn(parser, op, ...resolvedArgs);
if (result === null) {
shouldParse = false;
} else {
shouldParse = 'skip_first';
}
}
} else if (firstElem.type === 'qualified') {
if (firstElem.args && firstElem.args.length > 0) {
const [arg] = firstElem.args;
if (arg.startsWith('$')) {
shouldParse = peek(_.Token.hash_identifier);
} else if (arg.startsWith('type($')) {
shouldParse = peek(_.Token.exclamation_identifier, _.Token.bare_identifier);
}
}
}
}
if (shouldParse) {
// Recursively parse nested elements using the same parseDirective method
// If shouldParse === 'skip_first', the anchor element was already consumed
const firstElemIdx = shouldParse === 'skip_first' ? directive.elements.indexOf(firstElem) + 1 : 0;
for (let elemIdx = firstElemIdx; elemIdx < directive.elements.length; elemIdx++) {
this.parseDirective(directive.elements[elemIdx], parser, op, opInfo, directive.elements, elemIdx, vars);
}
}
break;
}
case 'conditional_alternative': {
const tryParseLiteral = (value) => {
switch (value) {
case '(': return parser.parseOptionalLParen();
case ')': return parser.parseOptionalRParen();
case '[': return parser.parseOptionalLSquare();
case ']': return parser.parseOptionalRSquare();
case '{': return parser.parseOptionalLBrace();
case '}': return parser.parseOptionalRBrace();
case ':': return parser.parseOptionalColon();
case '=': return parser.parseOptionalEqual();
case ',': return parser.parseOptionalComma();
case '<': return parser.parseOptionalLess();
case '>': return parser.parseOptionalGreater();
case '?': return parser.parseOptionalQuestion();
case '->': return parser.parseOptionalArrow();
case '...': return parser.parseOptionalEllipsis();
default: return parser.parseOptionalKeyword(value);
}
};
const checkMatch = (elem) => {
if (elem.type === 'literal') {
return tryParseLiteral(elem.value);
}
if (elem.type === 'operand_ref') {
return peek(_.Token.percent_identifier);
}
if (elem.type === 'attribute_ref') {
return peek(_.Token.bare_identifier, _.Token.integer, _.Token.floatliteral, _.Token.l_square, _.Token.at_identifier, _.Token.hash_identifier);
}
if (elem.type === 'region_ref') {
return peek(_.Token.l_brace);
}
if (elem.type === 'successor_ref') {
return peek(_.Token.caret_identifier);
}
if (elem.type === 'custom') {
return peek(_.Token.bare_identifier, _.Token.integer, _.Token.minus, _.Token.percent_identifier, _.Token.l_square, _.Token.l_paren, _.Token.question);
}
return false;
};
const firstElem = directive.firstAlt.find((e) => e.type !== 'whitespace');
let matchedFirst = firstElem && checkMatch(firstElem);
let customDirectiveHandledFirst = false;
// For custom directives, try calling them and check if they return null
if (matchedFirst && firstElem.type === 'custom') {
const fn = this._customDirectives.get(firstElem.parser);
if (fn) {
const resolvedArgs = (firstElem.args || []).map((arg) => {
if (typeof arg === 'string' && arg.startsWith('$')) {
return arg.slice(1); // Strip $ prefix to get attribute name
}
return arg;
});
const result = fn(parser, op, ...resolvedArgs);
if (result === null) {
matchedFirst = false;
} else {
customDirectiveHandledFirst = true;
}
}
}
const literalHandledFirst = matchedFirst && firstElem.type === 'literal';
if (matchedFirst) {
const skipFirst = customDirectiveHandledFirst || literalHandledFirst;
const startIdx = skipFirst ? directive.firstAlt.indexOf(firstElem) + 1 : 0;
for (let elemIdx = startIdx; elemIdx < directive.firstAlt.length; elemIdx++) {
this.parseDirective(directive.firstAlt[elemIdx], parser, op, opInfo, directive.firstAlt, elemIdx, vars);
}
} else if (directive.secondOptional) {
const secondElem = directive.secondAlt.find((e) => e.type !== 'whitespace');
const matchedSecond = secondElem && checkMatch(secondElem);
if (matchedSecond) {
const startIdx = secondElem.type === 'literal' ? directive.secondAlt.indexOf(secondElem) + 1 : 0;
for (let elemIdx = startIdx; elemIdx < directive.secondAlt.length; elemIdx++) {
this.parseDirective(directive.secondAlt[elemIdx], parser, op, opInfo, directive.secondAlt, elemIdx, vars);
}
}
} else if (directive.secondAlt && directive.secondAlt.length > 0) {
for (const elem of directive.secondAlt) {
this.parseDirective(elem, parser, op, opInfo, directive.secondAlt, 0, vars);
}
}
break;
}
default: {
parser.emitError(`Unsupported directive type '${directive.type}'`);
}
}
}
parseType(parser, dialect) {
return new _.Type(`!${dialect}.${parser.getFullSymbolSpec()}`);
}
parseAttribute(/* parser, type */) {
return null;
}
parseOperation(parser, result) {
const opInfo = result.name.getRegisteredInfo();
if (!opInfo) {
return false;
}
if ((opInfo.metadata.hasParser || opInfo.metadata.hasCustomAssemblyFormat) && !opInfo.metadata.assemblyFormat) {
const op = result.name.getStringRef();
parser.emitError(`Operation parser '${op}' not implemented`);
}
// Mark as using assembly format parsing (bypasses validation check)
if (result.compatibility === undefined && opInfo.metadata.assemblyFormat) {
result.compatibility = true;
}
const vars = new Map();
for (const input of opInfo.metadata?.operands || []) {
vars.set(input.name, { operands: [], types: [] });
}
for (const resultInfo of opInfo.metadata?.results || []) {
vars.set(resultInfo.name, { types: [] });
}
const directives = opInfo.directives || [];
for (let i = 0; i < directives.length; i++) {
this.parseDirective(directives[i], parser, result, opInfo, directives, i, vars);
}
for (const [, v] of vars) {
if (v.operands?.length > 0 && v.types?.length > 0) {
parser.resolveOperands(v.operands, v.types, result.operands);
} else if (v.operands?.length > 0 && result.types.length > 0) {
// SameOperandsAndResultType: use result type for operands
const types = v.operands.map(() => result.types[0]);
parser.resolveOperands(v.operands, types, result.operands);
} else if (v.operands?.length > 0) {
// No explicit type - resolve from scope (type from definition)
for (const operand of v.operands) {
parser.resolveOperand(operand, null, result.operands);
}
}
}
// AttrSizedOperandSegments trait
if (vars && opInfo.metadata?.operands && result.name.hasTrait('AttrSizedOperandSegments')) {
const segmentSizes = [];
for (const operandMeta of opInfo.metadata.operands) {
const entry = vars.get(operandMeta.name);
segmentSizes.push(entry?.operands?.length || 0);
}
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr(segmentSizes));
}
this.inferResultTypes(result, vars);
return true;
}
parseCustomTypeWithFallback(parser, type) {
const optionalType = parser.parseOptionalType();
if (optionalType) {
return optionalType;
}
if (type && this._customTypes.has(type.name)) {
let typeT = this._customTypes.get(type.name);
if (typeof typeT !== 'function') {
typeT = { type, name: typeT };
}
return parser.parseCustomTypeWithFallback(typeT);
}
throw new mlir.Error(`Unsupported type constraint '${type}'.`);
}
parseCustomAttributeWithFallback(parser, type) {
if (type && this._customAttributes.has(type.name)) {
const attrT = this._customAttributes.get(type.name);
return parser.parseCustomAttributeWithFallback(attrT, type);
}
if (type && Array.isArray(type.values)) {
const value = parser.parseOptionalKeyword(type.values);
if (value !== null) {
return new _.TypedAttr(value, null);
}
}
return parser.parseOptionalAttribute(null);
}
parseOptionalAttr(parser, type) {
if (!Array.isArray(type.args) || type.args.length === 0) {
throw new mlir.Error(`Invalid 'OptionalAttr' type.`);
}
const [elementType] = type.args;
return this.parseCustomAttributeWithFallback(parser, elementType);
}
parseDefaultValuedAttr(parser, type) {
if (!Array.isArray(type.args) || type.args.length === 0) {
throw new mlir.Error(`Invalid 'DefaultValuedAttr' type.`);
}
const [elementType] = type.args;
return this.parseCustomAttributeWithFallback(parser, elementType);
}
parseDefaultValuedOptionalAttr(parser, type) {
if (!Array.isArray(type.args) || type.args.length === 0) {
throw new mlir.Error(`Invalid 'DefaultValuedOptionalAttr' type.`);
}
const [elementType] = type.args;
return this.parseCustomAttributeWithFallback(parser, elementType);
}
parseTypeAttrOf(parser, type) {
if (!Array.isArray(type.args) || type.args.length === 0) {
throw new mlir.Error(`Invalid 'TypeAttrOf' type.`);
}
const parsedType = parser.parseOptionalType();
if (parsedType) {
return { value: parsedType, type: 'type' };
}
return null;
}
parseAnyAttrOf(parser) {
// This allows parseAttribute to handle the full syntax including `: type` suffix
return parser.parseOptionalAttribute(null);
}
parseArrayAttr(parser) {
// Try parsing array attribute (starts with '[') or attribute alias reference (starts with '#')
const attr = parser.parseOptionalAttribute();
if (attr) {
return attr;
}
return null;
}
parseConfinedAttr(parser, type) {
if (!Array.isArray(type.args) || type.args.length === 0) {
throw new mlir.Error(`Invalid ConfinedAttr type.`);
}
const [baseType] = type.args;
return this.parseCustomAttributeWithFallback(parser, baseType);
}
parseTypedAttrInterface(parser) {
return parser.parseAttribute();
}
parseUnitAttr(parser) {
parser.parseOptionalKeyword('unit');
return new _.UnitAttr();
}
parseEnumAttr(parser, type) {
const value = parser.parseOptionalKeyword(type.values);
if (value !== null) {
return new _.TypedAttr(value, null);
}
return null;
}
parseSymbolNameAttr(parser) {
const value = parser.parseOptionalSymbolName();
if (value) {
return new _.StringAttr(value);
}
return null;
}
parseSymbolRefAttr(parser) {
return parser.parseAttribute();
}
parseFlatSymbolRefAttr(parser) {
return this.parseSymbolRefAttr(parser);
}
// custom<DynamicIndexList>($dynamic_operands, $static_attr, $scalable_attr?, "Delimiter::Paren"?)
parseDynamicIndexList(parser, op, operandsAttr, staticAttrName, scalableAttrName, delimiterSpec) {
// Determine delimiter from delimiterSpec
let openDelim = '[';
if (typeof delimiterSpec === 'string' && delimiterSpec.includes('Paren')) {
openDelim = '(';
}
const unresolvedOperands = [];
const staticValues = [];
const scalableFlags = [];
const opened = openDelim === '(' ? parser.parseOptionalLParen() : parser.parseOptionalLSquare();
if (opened) {
const closed = () => openDelim === '(' ? parser.parseOptionalRParen() : parser.parseOptionalRSquare();
if (!closed()) {
do {
const isScalable = parser.parseOptionalLSquare();
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedOperands.push(operand);
staticValues.push(_.ShapedType.kDynamic);
if (parser.parseOptionalColon()) {
parser.parseType();
}
} else {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
staticValues.push(intVal);
}
scalableFlags.push(isScalable);
if (isScalable) {
if (!parser.parseOptionalRSquare()) {
parser.emitError(`Expected ']' for scalable index`);
}
}
} while (parser.parseOptionalComma());
if (openDelim === '(') {
parser.parseRParen();
} else {
parser.parseRSquare();
}
}
}
const indexType = new _.IndexType();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => indexType), op.operands);
if (staticAttrName && staticValues.length > 0) {
op.addAttribute(staticAttrName, staticValues);
}
if (scalableFlags.length > 0 && scalableFlags.some((f) => f)) {
if (!scalableAttrName) {
parser.emitError(`Scalable indices found but no scalable attribute name provided`);
}
op.addAttribute(scalableAttrName, scalableFlags);
}
}
parseNestedDynamicIndexList(parser, op, operandsArr, staticAttrName) {
const unresolvedOperands = [];
const innerArrayAttrs = [];
parser.parseLSquare();
do {
const staticVals = [];
parser.parseLSquare();
do {
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedOperands.push(operand);
staticVals.push(_.ShapedType.kDynamic);
} else {
staticVals.push(parser.parseInteger());
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
innerArrayAttrs.push(staticVals);
} while (parser.parseOptionalComma());
parser.parseRSquare();
const indexType = new _.IndexType();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => indexType), op.operands);
if (staticAttrName && innerArrayAttrs.length > 0) {
op.addAttribute(staticAttrName, innerArrayAttrs);
}
}
parseOffsets(parser, op, attrName) {
const values = [];
for (;;) {
if (parser.parseOptionalMinus()) {
values.push(-parser.parseInteger());
} else {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
values.push(intVal);
}
if (!parser.parseOptionalComma()) {
break;
}
}
if (attrName) {
op.addAttribute(attrName, values);
}
}
parseSymbolVisibility(parser, op) {
const visibility = parser.parseOptionalKeyword(['public', 'private', 'nested']);
if (visibility) {
op.addAttribute('sym_visibility', visibility);
}
}
parseTypeOrAttr(parser, op, typeArg, attrArg) {
if (parser.parseOptionalEqual()) {
const attr = parser.parseAttribute();
if (parser.parseOptionalColon()) {
const type = parser.parseType();
op.addAttribute(typeArg, type);
attr.type = type;
} else if (attr && attr.type) {
op.addAttribute(typeArg, attr.type);
}
op.addAttribute(attrArg, attr);
return;
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
op.addAttribute(typeArg, type);
if (parser.parseOptionalEqual()) {
const attr = parser.parseAttribute();
if (parser.parseOptionalColon()) {
const attrType = parser.parseType();
attr.type = attrType;
}
op.addAttribute(attrArg, attr);
}
return;
}
parser.emitError("Expected ':' or '=' in TypeOrAttr");
}
parseSizeAwareType(parser, op, typeArg) {
const type = parser.parseType();
parser.parseLBrace();
const operand = parser.parseOperand();
parser.parseRBrace();
if (!Array.isArray(typeArg)) {
throw new mlir.Error(`Invalid argument 'typeArg'.`);
}
if (typeArg.length === 0) {
typeArg.push(type);
} else {
typeArg[0] = type;
}
if (operand) {
parser.resolveOperand(operand, null, op.operands);
}
}
parseCopyOpRegion(parser, op /*, regionName, inputTypes, outputTypes */) {
op.regions.push({ blocks: [] });
}
parseResultTypeList(parser, op) {
const types = parser.parseCommaSeparatedList('paren', () => parser.parseType());
if (types.length > 0) {
op.addAttribute('result_types', types);
}
}
parsePipelineAttr(parser, op, attrName) {
const keyword = parser.parseOptionalKeyword();
if (keyword) {
op.addAttribute(attrName, keyword);
} else {
const attr = parser.parseAttribute();
op.addAttribute(attrName, attr);
}
}
parseKnobsDictionary(parser, op, attrName) {
parser.parseKeyword('knobs');
parser.parseEqual();
const attr = parser.parseAttribute();
op.addAttribute(attrName, attr);
}
parseCmdParameterGatherOperations(parser, op, sourceScopeOps, sourceKeysOps, sourceOffsetsOps, targetOps, targetTypes, targetSizeOps, targetOffsetsOps, targetLengthsOps) {
do {
const firstOperand = parser.parseOperand();
let key = firstOperand;
if (parser.parseOptionalColon()) {
sourceScopeOps.push(firstOperand);
parser.parseColon();
key = parser.parseOperand();
}
sourceKeysOps.push(key);
parser.parseLSquare();
sourceOffsetsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseArrow();
const rowTarget = parser.parseOperand();
if (targetOps.length === 0) {
targetOps.push(rowTarget);
}
parser.parseLSquare();
targetOffsetsOps.push(parser.parseOperand());
parser.parseKeyword('for');
targetLengthsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseColon();
const rowType = parser.parseType();
if (targetTypes.length === 0) {
targetTypes.push(rowType);
}
parser.parseLBrace();
const rowSize = parser.parseOperand();
if (targetSizeOps.length === 0) {
targetSizeOps.push(rowSize);
}
parser.parseRBrace();
} while (parser.parseOptionalComma());
}
// custom<AsyncParameterGatherOperations>(...)
// Parses: %scope::%key[%offset] -> %target[%offset to %end for %length] : type{%size}, ...
parseAsyncParameterGatherOperations(parser, op, sourceScopeOps, sourceKeysOps, sourceOffsetsOps, targetOps, targetTypes, targetSizeOps, targetOffsetsOps, targetEndsOps, targetLengthsOps) {
do {
const firstOperand = parser.parseOperand();
let key = firstOperand;
if (parser.parseOptionalColon()) {
sourceScopeOps.push(firstOperand);
parser.parseColon();
key = parser.parseOperand();
}
sourceKeysOps.push(key);
parser.parseLSquare();
sourceOffsetsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseArrow();
const rowTarget = parser.parseOperand();
if (targetOps.length === 0) {
targetOps.push(rowTarget);
}
parser.parseLSquare();
targetOffsetsOps.push(parser.parseOperand());
parser.parseKeyword('to');
targetEndsOps.push(parser.parseOperand());
parser.parseKeyword('for');
targetLengthsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseColon();
const rowType = parser.parseType();
if (targetTypes.length === 0) {
targetTypes.push(rowType);
}
parser.parseLBrace();
const rowSize = parser.parseOperand();
if (targetSizeOps.length === 0) {
targetSizeOps.push(rowSize);
}
parser.parseRBrace();
} while (parser.parseOptionalComma());
}
// custom<CmdParameterScatterOperations>(...)
// Parses: %source[%offset for %length] : type{%size} -> %scope::%key[%offset], ...
parseCmdParameterScatterOperations(parser, op, sourceOps, sourceTypes, sourceSizeOps, sourceOffsetsOps, sourceLengthsOps, targetScopeOps, targetKeysOps, targetOffsetsOps) {
do {
const rowSource = parser.parseOperand();
if (sourceOps.length === 0) {
sourceOps.push(rowSource);
}
parser.parseLSquare();
sourceOffsetsOps.push(parser.parseOperand());
parser.parseKeyword('for');
sourceLengthsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseColon();
const rowType = parser.parseType();
if (sourceTypes.length === 0) {
sourceTypes.push(rowType);
}
parser.parseLBrace();
const rowSize = parser.parseOperand();
if (sourceSizeOps.length === 0) {
sourceSizeOps.push(rowSize);
}
parser.parseRBrace();
// Parse -> %scope::%key[%offset]
parser.parseArrow();
const firstOperand = parser.parseOperand();
let key = firstOperand;
if (parser.parseOptionalColon()) {
targetScopeOps.push(firstOperand);
parser.parseColon();
key = parser.parseOperand();
}
targetKeysOps.push(key);
parser.parseLSquare();
targetOffsetsOps.push(parser.parseOperand());
parser.parseRSquare();
} while (parser.parseOptionalComma());
}
// custom<AsyncParameterScatterOperations>(...)
// Parses: %source[%offset to %end for %length] : type{%size} -> %scope::%key[%offset], ...
parseAsyncParameterScatterOperations(parser, op, sourceOps, sourceTypes, sourceSizeOps, sourceOffsetsOps, sourceEndsOps, sourceLengthsOps, targetScopeOps, targetKeysOps, targetOffsetsOps) {
do {
const rowSource = parser.parseOperand();
if (sourceOps.length === 0) {
sourceOps.push(rowSource);
}
parser.parseLSquare();
sourceOffsetsOps.push(parser.parseOperand());
parser.parseKeyword('to');
sourceEndsOps.push(parser.parseOperand());
parser.parseKeyword('for');
sourceLengthsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseColon();
const rowType = parser.parseType();
if (sourceTypes.length === 0) {
sourceTypes.push(rowType);
}
parser.parseLBrace();
const rowSize = parser.parseOperand();
if (sourceSizeOps.length === 0) {
sourceSizeOps.push(rowSize);
}
parser.parseRBrace();
// Parse -> %scope::%key[%offset]
parser.parseArrow();
const firstOperand = parser.parseOperand();
let key = firstOperand;
if (parser.parseOptionalColon()) {
targetScopeOps.push(firstOperand);
parser.parseColon();
key = parser.parseOperand();
}
targetKeysOps.push(key);
parser.parseLSquare();
targetOffsetsOps.push(parser.parseOperand());
parser.parseRSquare();
} while (parser.parseOptionalComma());
}
parseEnumFlags(parser, type, separator, optional) {
const flags = [];
do {
if (optional && flags.length === 0) {
const value = parser.parseOptionalKeyword(type.values);
if (!value) {
return null;
}
flags.push(value);
} else {
const value = parser.parseKeyword();
if (!type.values.includes(value)) {
parser.emitError(`Invalid enum value '${value}'`);
}
flags.push(value);
}
} while (separator === ',' ? parser.parseOptionalComma() : parser.parseOptionalVerticalBar());
return new _.TypedAttr(flags.join(', '));
}
parseEnumFlagsAngleBracketComma(parser, type) {
if (parser.parseOptionalLess()) {
const value = this.parseEnumFlags(parser, type, ',');
parser.parseGreater();
return value;
}
return parser.parseOptionalAttribute();
}
parseEnumFlagsAngleBracketPipe(parser, type) {
if (parser.parseOptionalLess()) {
const value = this.parseEnumFlags(parser, type, '|');
parser.parseGreater();
return value;
}
return parser.parseOptionalAttribute();
}
parseEnumFlagsPipe(parser, type) {
const value = this.parseEnumFlags(parser, type, '|', true);
if (value === null) {
return parser.parseOptionalAttribute();
}
return value;
}
parseOptional(parser) {
return parser.parseOptionalType();
}
};
_.hlo = {};
_.hlo.HLODialect = class extends _.Dialect {
constructor(operations, name) {
super(operations, name);
this.registerCustomDirective('SameOperandsAndResultType', this.parseSameOperandsAndResultType.bind(this));
this.registerCustomDirective('VariadicSameOperandsAndResultType', this.parseVariadicSameOperandsAndResultType.bind(this));
this.registerCustomDirective('ComplexOpType', this.parseComplexOpType.bind(this));
this.registerCustomDirective('SelectOpType', this.parseSelectOpType.bind(this));
this.registerCustomDirective('TupleOpType', this.parseTupleOpType.bind(this));
this.registerCustomDirective('PairwiseOpType', this.parsePairwiseOpType.bind(this));
this.registerCustomDirective('ConvolutionDimensions', this.parseConvolutionDimensions.bind(this));
this.registerCustomDirective('DotDimensionNumbers', this.parseDotDimensionNumbers.bind(this));
this.registerCustomDirective('PrecisionConfig', this.parsePrecisionConfig.bind(this));
this.registerCustomDirective('PrecisionConfigAndAlgorithm', this.parsePrecisionConfigAndAlgorithm.bind(this));
this.registerCustomDirective('WindowAttributes', this.parseWindowAttributes.bind(this));
this.registerCustomDirective('SliceRanges', this.parseSliceRanges.bind(this));
this.registerCustomDirective('CustomCallTarget', this.parseCustomCallTarget.bind(this));
this.registerCustomDirective('VariadicOperandWithAttribute', this.parseVariadicOperandWithAttribute.bind(this));
}
assignFromFunctionType(operandArrays, resultArray, fnType) {
for (let i = 0; i < operandArrays.length && i < fnType.inputs.length; i++) {
operandArrays[i].push(fnType.inputs[i]);
}
if (fnType.results.length === 1) {
resultArray.push(fnType.results[0]);
}
}
parseSameOperandsAndResultType(parser, op, ...typeArrays) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
const operandArrays = typeArrays.slice(0, -1);
const resultArray = typeArrays[typeArrays.length - 1];
this.assignFromFunctionType(operandArrays, resultArray, type);
} else {
for (const arr of typeArrays) {
arr.push(type);
}
}
}
// custom<VariadicSameOperandsAndResultType>(ref($inputs), type($inputs), type($result))
parseVariadicSameOperandsAndResultType(parser, op, operands, operandTypes, resultTypes) {
const type = parser.parseType();
// All operands get the same type
for (let i = 0; i < operands.length; i++) {
operandTypes.push(type);
}
// Result also gets the same type
resultTypes.push(type);
}
// custom<ComplexOpType>(ref($operands), type($operands), type($result))
parseComplexOpType(parser, op, operands, operandTypes, resultTypes) {
const type = parser.parseType();
for (let i = 0; i < operands.length; i++) {
operandTypes.push(type);
}
resultTypes.push(type);
}
// custom<SelectOpType>(type($pred), type($on_true), type($on_false), type($result))
parseSelectOpType(parser, op, predTypes, onTrueTypes, onFalseTypes, resultTypes) {
const firstType = parser.parseType();
if (parser.parseOptionalComma()) {
const secondType = parser.parseType();
predTypes.push(firstType);
onTrueTypes.push(secondType);
onFalseTypes.push(secondType);
resultTypes.push(secondType);
} else {
predTypes.push(firstType);
onTrueTypes.push(firstType);
onFalseTypes.push(firstType);
resultTypes.push(firstType);
}
}
// custom<TupleOpType>(type($operands), type($result))
parseTupleOpType(parser, op, operandTypes, resultTypes) {
const type = parser.parseType();
operandTypes.push(type);
resultTypes.push(type);
}
// custom<PairwiseOpType>(type($operands), type($results))
parsePairwiseOpType(parser, op, operandTypes, resultTypes) {
while (true) {
const type = parser.parseType();
if (!type) {
break;
}
operandTypes.push(type);
resultTypes.push(type);
if (!parser.parseOptionalComma()) {
break;
}
}
}
parseDims(parser) {
const dims = [];
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
const kw = parser.parseOptionalKeyword();
if (kw) {
dims.push(kw);
} else {
break;
}
} else {
dims.push(intVal);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
return dims;
}
parseConvolutionDimensions(parser, op, attrName) {
const input = this.parseDims(parser);
parser.parseKeyword('x');
const kernel = this.parseDims(parser);
parser.parseArrow();
const output = this.parseDims(parser);
op.addAttribute(attrName, new _.ConvDimensionNumbersAttr(input, kernel, output));
}
parseWindowAttributes(parser, op, stridesAttr, paddingAttr, lhsDilationAttr, rhsDilationAttr, reversalAttr) {
const windowAttrs = {
stride: [],
pad: [],
lhs_dilate: [],
rhs_dilate: [],
window_reversal: []
};
const parseArray = () => {
return parser.parseCommaSeparatedList('square', () => {
const intValue = parser.parseOptionalInteger();
if (intValue !== null) {
return intValue;
}
const kw = parser.parseOptionalKeyword();
if (kw) {
return kw;
}
// Nested array
return parseArray();
});
};
for (;;) {
const key = parser.parseOptionalKeyword();
if (key) {
if (parser.parseOptionalEqual()) {
windowAttrs[key] = parseArray();
}
parser.parseOptionalComma();
} else {
break;
}
}
if (stridesAttr && windowAttrs.stride.length > 0) {
op.addAttribute(stridesAttr, windowAttrs.stride);
}
if (paddingAttr && windowAttrs.pad.length > 0) {
op.addAttribute(paddingAttr, windowAttrs.pad);
}
if (lhsDilationAttr && windowAttrs.lhs_dilate.length > 0) {
op.addAttribute(lhsDilationAttr, windowAttrs.lhs_dilate);
}
if (rhsDilationAttr && windowAttrs.rhs_dilate.length > 0) {
op.addAttribute(rhsDilationAttr, windowAttrs.rhs_dilate);
}
if (reversalAttr && windowAttrs.window_reversal.length > 0) {
op.addAttribute(reversalAttr, windowAttrs.window_reversal);
}
}
parseDotDimensionNumbers(parser, op, attrName = 'dot_dimension_numbers') {
const dimensions = {
lhs_batching_dimensions: [],
rhs_batching_dimensions: [],
lhs_contracting_dimensions: [],
rhs_contracting_dimensions: []
};
const parseIntArray = () => {
return parser.parseCommaSeparatedList('optionalSquare', () => {
const value = parser.parseOptionalInteger();
if (value !== null) {
return value;
}
parser.parseKeyword();
return null;
});
};
const parsePair = () => {
const first = parseIntArray();
let second = [];
if (parser.parseOptionalKeyword('x')) {
second = parseIntArray();
}
return { first, second };
};
if (parser.parseOptionalKeyword('batching_dims') || parser.parseOptionalKeyword('batch_dims')) {
parser.parseEqual();
const pair = parsePair();
dimensions.lhs_batching_dimensions = pair.first;
dimensions.rhs_batching_dimensions = pair.second;
parser.parseOptionalComma();
}
parser.parseKeyword('contracting_dims');
parser.parseEqual();
const pair = parsePair();
dimensions.lhs_contracting_dimensions = pair.first;
dimensions.rhs_contracting_dimensions = pair.second;
op.addAttribute(attrName, dimensions);
}
parsePrecisionConfig(parser, op /*, args */) {
parser.parseOptionalComma();
if (!parser.parseOptionalKeyword('precision')) {
return;
}
parser.parseEqual();
const precision = parser.parseCommaSeparatedList('square', () => {
const kw = parser.parseOptionalKeyword();
if (kw) {
return kw;
}
parser.parseKeyword();
return null;
});
if (precision.length > 0) {
op.addAttribute('precision_config', precision);
}
}
parsePrecisionConfigAndAlgorithm(parser, op /*, args */) {
if (!parser.parseOptionalComma()) {
return;
}
if (parser.parseOptionalKeyword('algorithm')) {
parser.parseOptionalEqual();
const algorithm = parser.parseAttribute();
op.addAttribute('algorithm', algorithm);
return;
}
if (parser.parseOptionalKeyword('precision')) {
parser.parseOptionalEqual();
const precision = parser.parseCommaSeparatedList('optionalSquare', () => {
const kw = parser.parseOptionalKeyword();
if (kw) {
return kw;
}
parser.parseKeyword();
return null;
});
if (precision.length > 0) {
op.addAttribute('precision_config', precision);
}
if (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('algorithm')) {
parser.parseOptionalEqual();
const algorithm = parser.parseAttribute();
op.addAttribute('algorithm', algorithm);
}
}
}
}
parseSliceRanges(parser, op /*, args */) {
const ranges = {
start_indices: [],
limit_indices: [],
strides: []
};
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
let intVal = parser.parseOptionalInteger();
if (intVal !== null) {
ranges.start_indices.push(intVal);
}
parser.parseOptionalColon();
intVal = parser.parseOptionalInteger();
if (intVal !== null) {
ranges.limit_indices.push(intVal);
}
if (parser.parseOptionalColon()) {
intVal = parser.parseOptionalInteger();
if (intVal !== null) {
ranges.strides.push(intVal);
}
} else {
ranges.strides.push(1);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
op.addAttribute('start_indices', ranges.start_indices);
op.addAttribute('limit_indices', ranges.limit_indices);
op.addAttribute('strides', ranges.strides);
}
// custom<CustomCallTarget>($call_target_name)
parseCustomCallTarget(parser, op, attrName) {
let target = null;
const sym = parser.parseOptionalSymbolName();
if (sym) {
target = `@${sym.value || sym}`;
} else {
const str = parser.parseOptionalString();
if (str === null) {
parser.emitError("Expected '@' or string for CustomCallTarget");
}
target = str;
}
op.addAttribute(attrName || 'call_target_name', target);
}
// custom<VariadicOperandWithAttribute>($inputs)
parseVariadicOperandWithAttribute(parser, op, operands) {
let operand = parser.parseOptionalOperand();
while (operand) {
const attrs = new Map();
parser.parseOptionalAttrDict(attrs);
if (attrs.size > 0) {
operand.attributes = attrs;
}
operands.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
}
parseReduceOp(parser, result, createDimensions, returnOpName, lparenConsumed = false) {
const unresolvedOperands = [];
const unresolvedInitOperands = [];
while (true) {
if (lparenConsumed) {
lparenConsumed = false;
} else if (!parser.parseOptionalLParen()) {
break;
}
const operand = parser.parseOperand();
parser.parseKeyword('init');
parser.parseColon();
const initOperand = parser.parseOperand();
parser.parseRParen();
unresolvedOperands.push(operand);
unresolvedInitOperands.push(initOperand);
parser.parseOptionalComma();
}
const allUnresolved = unresolvedOperands.concat(unresolvedInitOperands);
if (parser.parseOptionalKeyword('applies')) {
const innerOpName = parser.parseCustomOperationName();
parser.parseKeyword('across');
parser.parseKeyword('dimensions');
parser.parseEqual();
const dimensions = parser.parseCommaSeparatedList('square', () => {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
parser.emitError('Expected integer dimension in reduce operation');
}
return intVal;
});
result.addAttribute('dimensions', createDimensions(dimensions));
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(allUnresolved, type.inputs, result.operands);
result.addTypes(type.results);
} else {
const types = Array.isArray(type) ? type : [type];
parser.resolveOperands(allUnresolved, allUnresolved.map((_, i) => types[i] || types[0]), result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
}
} else {
parser.resolveOperands(allUnresolved, allUnresolved.map(() => null), result.operands);
}
const region = { blocks: [] };
const block = { operations: [], arguments: [] };
let tensorType = null;
if (result.operands.length > 0 && result.operands[0].type) {
const operandType = result.operands[0].type;
if (operandType instanceof _.RankedTensorType || operandType instanceof _.UnrankedTensorType) {
tensorType = new _.RankedTensorType([], operandType.elementType);
}
}
block.arguments.push({ value: '%lhs', type: tensorType });
block.arguments.push({ value: '%rhs', type: tensorType });
const innerOp = new _.OperationState(null, innerOpName);
innerOp.operands.push(new _.Value('%lhs', tensorType));
innerOp.operands.push(new _.Value('%rhs', tensorType));
if (tensorType) {
innerOp.addTypes([tensorType]);
}
block.operations.push(_.Operation.create(innerOp));
const returnOp = new _.OperationState(null, this.getOperation(returnOpName));
returnOp.operands.push(new _.Value('%0', tensorType));
block.operations.push(_.Operation.create(returnOp));
region.blocks.push(block);
result.regions.push(region);
return true;
}
// Non-compact syntax: parse "across dimensions = [...] : type reducer (args) { ... }"
parser.parseKeyword('across');
parser.parseKeyword('dimensions');
parser.parseEqual();
const dimensions = parser.parseCommaSeparatedList('square', () => {
return parser.parseInteger();
});
result.addAttribute('dimensions', createDimensions(dimensions));
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(allUnresolved, type.inputs, result.operands);
result.addTypes(type.results);
} else {
const types = Array.isArray(type) ? type : [type];
parser.resolveOperands(allUnresolved, types, result.operands);
}
}
parser.parseKeyword('reducer');
// Parse block arguments: (%lhs : type, %rhs : type) (%linit : type, %rinit : type)
const regionArgs = [];
while (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const value = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
regionArgs.push({ value, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
return true;
}
// Parse scan operation format: scan (%inputs) inits (%inits) dimension=0 { body } : types
parseScanOp(parser, result, returnOpName, lparenConsumed = false) {
let unresolvedInputs = null;
if (lparenConsumed) {
unresolvedInputs = [];
if (!parser.parseOptionalRParen()) {
do {
unresolvedInputs.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
unresolvedInputs = parser.parseCommaSeparatedList('paren', () => parser.parseOperand());
}
parser.parseKeyword('inits');
const unresolvedInits = parser.parseCommaSeparatedList('paren', () => parser.parseOperand());
parser.parseKeyword('dimension');
parser.parseOptionalEqual();
result.addAttribute('dimension', parser.parseInteger());
// Parse optional attributes: is_reverse=true, is_associative=true
while (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('is_reverse')) {
parser.parseOptionalEqual();
const value = parser.parseKeyword();
result.addAttribute('is_reverse', value === 'true');
} else if (parser.parseOptionalKeyword('is_associative')) {
parser.parseOptionalEqual();
const value = parser.parseKeyword();
result.addAttribute('is_associative', value === 'true');
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region, undefined, true); // IsolatedFromAbove
// Parse types: : (input types, init types) -> (output types, carry types)
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
const allUnresolved = unresolvedInputs.concat(unresolvedInits);
parser.resolveOperands(allUnresolved, type.inputs, result.operands);
result.addTypes(type.results);
}
}
return true;
}
parseWhileOp(parser, result, lparenConsumed = false) {
if (!lparenConsumed) {
parser.parseLParen();
}
const unresolvedOperands = [];
if (!parser.parseOptionalRParen()) {
do {
parser.parseOperand(); // iter_arg
if (parser.parseOptionalEqual()) {
unresolvedOperands.push(parser.parseOperand());
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
const types = [];
if (unresolvedOperands.length > 0) {
parser.parseColon();
const typeList = parser.parseTypeList();
types.push(...typeList);
}
parser.resolveOperands(unresolvedOperands, types, result.operands);
result.addTypes(types);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseKeyword('cond');
const cond = result.addRegion();
parser.parseRegion(cond);
parser.parseKeyword('do');
const body = result.addRegion();
parser.parseRegion(body);
return true;
}
};
_.stablehlo = {};
_.stablehlo.FutureType = class extends _.Type {
constructor(types) {
super(null);
this.types = types;
}
static parse(parser) {
parser.parseLess();
const types = [];
do {
types.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.stablehlo.FutureType(types);
}
toString() {
return `!stablehlo.future<${this.types.map((t) => t.toString()).join(', ')}>`;
}
};
_.stablehlo.StableHLODialect = class extends _.hlo.HLODialect {
constructor(operations) {
super(operations, 'stablehlo');
this.registerCustomDirective('ExponentMantissa', this.parseExponentMantissa.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'token') {
return new _.Type(`!${dialect}.token`);
}
if (mnemonic === 'future') {
return _.stablehlo.FutureType.parse(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'stablehlo.constant') {
if (parser.parseOptionalLParen() && parser.parseOptionalRParen()) {
if (parser.parseOptionalLess()) {
result.propertiesAttr = parser.parseAttribute();
parser.parseGreater();
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
parser.parseLParen();
parser.parseRParen();
parser.parseArrow();
const type = parser.parseType();
result.addTypes([type]);
} else {
// Custom form: {attrs} value : type
parser.parseOptionalAttrDict(result.attributes);
const value = parser.parseAttribute();
if (value) {
result.addAttribute('value', value);
}
const types = parser.parseOptionalColonTypeList();
if (types.length > 0) {
result.addTypes([types[0]]);
} else if (value && value.type) {
result.addTypes([value.type]);
}
}
return true;
}
if (op === 'stablehlo.while') {
return super.parseWhileOp(parser, result);
}
if (op === 'stablehlo.reduce') {
return super.parseReduceOp(parser, result, (dims) => dims, 'stablehlo.return');
}
if (op === 'stablehlo.scan') {
return super.parseScanOp(parser, result, 'stablehlo.return');
}
return super.parseOperation(parser, result);
}
parseExponentMantissa(parser, op, exponentAttr, mantissaAttr) {
const keyword = parser.parseKeyword();
const match = /^e(\d+)m(\d+)$/.exec(keyword);
if (!match) {
throw new mlir.Error(`Expected exponent mantissa in format e#m#, got '${keyword}'`);
}
const exponent = parseInt(match[1], 10);
const mantissa = parseInt(match[2], 10);
op.addAttribute(exponentAttr, exponent);
op.addAttribute(mantissaAttr, mantissa);
}
};
_.vhlo = {};
_.vhlo.TypeV1Attr = class extends _.Attribute {
constructor(type) {
super();
this.type = type;
}
toString() {
return this.type.toString();
}
};
_.vhlo.VhloDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'vhlo');
this.registerCustomDirective('FunctionBody', this.parseFunctionBody.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'vhlo.constant_v1') {
parser.parseOptionalAttrDict(result.attributes);
const value = parser.parseAttribute();
if (value) {
result.addAttribute('value', value);
}
result.addTypes(parser.parseOptionalColonTypeList());
return true;
}
if (op === 'vhlo.return_v1') {
const unresolvedOperands = parser.parseOperandList();
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
return true;
}
return super.parseOperation(parser, result);
}
parseFunctionBody(parser, op /*, symNameAttr, bodyRegion, functionTypeAttr */) {
parser.parseFunctionOp(op, false);
}
readType(reader) {
const code = reader.readVarInt();
switch (code) {
case 0: return new _.IntegerType('i1'); // BooleanV1Type
case 1: return new _.ComplexType(reader.readType()); // ComplexV1Type
case 2: return new _.FloatType('bf16'); // FloatBF16V1Type
case 3: return new _.FloatType('f16'); // FloatF16V1Type
case 4: return new _.FloatType('f32'); // FloatF32V1Type
case 5: return new _.FloatType('f64'); // FloatF64V1Type
case 36: return new _.FloatType('f8E3M4'); // FloatF8E3M4V1Type
case 35: return new _.FloatType('f8E4M3'); // FloatF8E4M3V1Type
case 6: return new _.FloatType('f8E4M3FN'); // FloatF8E4M3FNV1Type
case 7: return new _.FloatType('f8E5M2'); // FloatF8E5M2V1Type
case 34: return new _.FloatType('tf32'); // FloatTF32V1Type
case 8: { // FunctionV1Type
const numInputs = reader.readVarInt();
const inputs = [];
for (let i = 0; i < numInputs; i++) {
inputs.push(reader.readType());
}
const numOutputs = reader.readVarInt();
const outputs = [];
for (let i = 0; i < numOutputs; i++) {
outputs.push(reader.readType());
}
return new _.FunctionType(inputs, outputs);
}
case 9: return new _.IndexType(); // IndexV1Type
case 31: return new _.IntegerType('si2'); // IntegerSI2V1Type
case 10: return new _.IntegerType('si4'); // IntegerSI4V1Type
case 11: return new _.IntegerType('si8'); // IntegerSI8V1Type
case 12: return new _.IntegerType('si16'); // IntegerSI16V1Type
case 13: return new _.IntegerType('si32'); // IntegerSI32V1Type
case 14: return new _.IntegerType('si64'); // IntegerSI64V1Type
case 32: return new _.IntegerType('ui2'); // IntegerUI2V1Type
case 15: return new _.IntegerType('ui4'); // IntegerUI4V1Type
case 16: return new _.IntegerType('ui8'); // IntegerUI8V1Type
case 17: return new _.IntegerType('ui16'); // IntegerUI16V1Type
case 18: return new _.IntegerType('ui32'); // IntegerUI32V1Type
case 19: return new _.IntegerType('ui64'); // IntegerUI64V1Type
case 20: { // RankedTensorV1Type
const shape = reader.readSignedVarInts().map(BigInt);
const elementType = reader.readType();
return new _.RankedTensorType(shape, elementType);
}
case 21: { // RankedTensorV1TypeWithEncoding
const encoding = reader.readAttribute();
const shape = reader.readSignedVarInts().map(BigInt);
const elementType = reader.readType();
return new _.RankedTensorType(shape, elementType, encoding);
}
case 22: return new _.Type('!vhlo.token'); // TokenV1Type
case 23: { // TupleV1Type
const numTypes = reader.readVarInt();
const types = [];
for (let i = 0; i < numTypes; i++) {
types.push(reader.readType());
}
return new _.TupleType(types);
}
case 24: { // UniformQuantizedV1Type
reader.readVarInt(); // flags
const storageType = reader.readType();
const expressedType = reader.readType();
const scale = reader.readAPFloatWithKnownSemantics(new _.FloatType('f64'));
const zeroPoint = reader.readSignedVarInt();
const storageTypeMin = reader.readSignedVarInt();
const storageTypeMax = reader.readSignedVarInt();
return new _.quant.UniformQuantizedType(storageType, [storageTypeMin, storageTypeMax], expressedType, scale, zeroPoint);
}
case 25: return new _.UnrankedTensorType(reader.readType()); // UnrankedTensorV1Type
case 26: return new _.Type('!vhlo.witness'); // WitnessV1Type
case 37: return new _.FloatType('f4E2M1FN'); // FloatF4E2M1FNV1Type
case 38: return new _.FloatType('f6E2M3FN'); // FloatF6E2M3FNV1Type
case 39: return new _.FloatType('f6E3M2FN'); // FloatF6E3M2FNV1Type
case 27: return new _.FloatType('f8E4M3FNUZ'); // FloatF8E4M3FNUZV1Type
case 28: return new _.FloatType('f8E5M2FNUZ'); // FloatF8E5M2FNUZV1Type
case 29: return new _.FloatType('f8E4M3B11FNUZ'); // FloatF8E4M3B11FNUZV1Type
case 40: return new _.FloatType('f8E8M0FNU'); // FloatF8E8M0FNUV1Type
case 30: { // UniformQuantizedPerAxisV1Type
reader.readVarInt(); // flags
const storageType = reader.readType();
const expressedType = reader.readType();
const quantizedDimension = reader.readVarInt();
const storageTypeMin = reader.readSignedVarInt();
const storageTypeMax = reader.readSignedVarInt();
const numScales = reader.readVarInt();
const scales = [];
for (let i = 0; i < numScales; i++) {
scales.push(reader.readAPFloatWithKnownSemantics(new _.FloatType('f64')));
}
const numZeroPoints = reader.readVarInt();
const zeroPoints = [];
for (let i = 0; i < numZeroPoints; i++) {
zeroPoints.push(reader.readSignedVarInt());
}
return new _.quant.UniformQuantizedPerAxisType(storageType, [storageTypeMin, storageTypeMax], expressedType, quantizedDimension, scales, zeroPoints);
}
case 33: return new _.NoneType(); // NoneV1Type
case 41: { // RankedBufferV1Type
const shape = reader.readSignedVarInts().map(BigInt);
const elementType = reader.readType();
return new _.RankedTensorType(shape, elementType);
}
case 42: { // FutureV1Type
const numTypes = reader.readVarInt();
const types = [];
for (let i = 0; i < numTypes; i++) {
types.push(reader.readType());
}
return new _.stablehlo.FutureType(types);
}
default:
throw new mlir.Error(`Unknown vhlo type code '${code}'.`);
}
}
readAttribute(reader) {
const code = reader.readVarInt();
switch (code) {
case 1: { // ArrayV1Attr
const count = reader.readVarInt();
const elements = [];
for (let i = 0; i < count; i++) {
elements.push(reader.readAttribute());
}
return new _.ArrayAttr(elements);
}
case 2: { // BooleanV1Attr
const value = reader.readVarInt();
return new _.BoolAttr(value === 1);
}
case 3: { // ComparisonDirectionV1Attr
const value = reader.readVarInt();
const directions = ['EQ', 'NE', 'GE', 'GT', 'LE', 'LT'];
return new _.StringAttr(directions[value] || `comparison_direction_${value}`);
}
case 4: { // ComparisonTypeV1Attr
const value = reader.readVarInt();
const types = ['NOTYPE', 'FLOAT', 'TOTALORDER', 'SIGNED', 'UNSIGNED'];
return new _.StringAttr(types[value] || `comparison_type_${value}`);
}
case 5: { // CustomCallApiVersionV1Attr
const value = reader.readVarInt();
return new _.StringAttr(`API_VERSION_${value}`);
}
case 6: { // DictionaryV1Attr
const count = reader.readVarInt();
const attrs = new Map();
for (let i = 0; i < count; i++) {
const name = reader.readAttribute();
const value = reader.readAttribute();
const key = name && name.value !== undefined ? name.value : `attr_${i}`;
attrs.set(key, value);
}
return { name: 'dictionary', value: attrs };
}
case 7: { // FftTypeV1Attr
const value = reader.readVarInt();
const types = ['FFT', 'IFFT', 'RFFT', 'IRFFT'];
return new _.StringAttr(types[value] || `fft_type_${value}`);
}
case 8: { // FloatV1Attr
const type = reader.readType();
const value = reader.readAPFloatWithKnownSemantics(type);
return new _.FloatAttr(type, value);
}
case 9: { // IntegerV1Attr
const type = reader.readType();
let bitWidth = 0;
if (type instanceof _.IndexType) {
bitWidth = 64;
} else {
const match = /(\d+)/.exec(type.name);
bitWidth = match ? parseInt(match[1], 10) : 64;
}
const value = reader.readAPIntWithKnownWidth(bitWidth);
return new _.IntegerAttr(type, value);
}
case 10: { // OutputOperandAliasV1Attr
const outputTupleIndices = reader.readSignedVarInts();
const operandIndex = reader.readSignedVarInt();
const operandTupleIndices = reader.readSignedVarInts();
return new _.StringAttr(`output_operand_alias<output_tuple_indices = [${outputTupleIndices}], operand_index = ${operandIndex}, operand_tuple_indices = [${operandTupleIndices}]>`);
}
case 11: { // PrecisionV1Attr
const value = reader.readVarInt();
const precisions = ['DEFAULT', 'HIGH', 'HIGHEST', 'PACKED_NIBBLE'];
return new _.StringAttr(precisions[value] || `precision_${value}`);
}
case 12: { // RngAlgorithmV1Attr
const value = reader.readVarInt();
const algs = ['DEFAULT', 'THREE_FRY', 'PHILOX'];
return new _.StringAttr(algs[value] || `rng_algorithm_${value}`);
}
case 13: { // RngDistributionV1Attr
const value = reader.readVarInt();
const dists = ['UNIFORM', 'NORMAL'];
return new _.StringAttr(dists[value] || `rng_distribution_${value}`);
}
case 14: { // StringV1Attr
const value = reader.readString();
return new _.StringAttr(value);
}
case 15: { // TensorV1Attr
const type = reader.readType();
const blob = reader.readBlob();
return new _.DenseElementsAttr(blob, type);
}
case 16: { // TransposeV1Attr
const value = reader.readVarInt();
const transposes = ['TRANSPOSE', 'NO_TRANSPOSE', 'ADJOINT'];
return new _.StringAttr(transposes[value] || `transpose_${value}`);
}
case 17: { // TypeV1Attr
const type = reader.readType();
return new _.vhlo.TypeV1Attr(type);
}
case 18: { // TypeExtensionsV1Attr
const bounds = reader.readSignedVarInts();
return new _.StringAttr(`type_extensions<bounds = [${bounds}]>`);
}
case 19: { // ResultAccuracyModeV1Attr
const value = reader.readVarInt();
const modes = ['DEFAULT', 'HIGH', 'HIGHEST', 'TOLERANCE'];
return new _.StringAttr(modes[value] || `result_accuracy_mode_${value}`);
}
case 20: { // ResultAccuracyV1Attr
const atol = reader.readAPFloatWithKnownSemantics(new _.FloatType('f64'));
const rtol = reader.readAPFloatWithKnownSemantics(new _.FloatType('f64'));
const ulps = reader.readSignedVarInt();
const mode = reader.readAttribute();
return new _.StringAttr(`result_accuracy<atol = ${atol}, rtol = ${rtol}, ulps = ${ulps}, mode = ${mode}>`);
}
default:
throw new mlir.Error(`Unknown vhlo attribute code '${code}'.`);
}
}
};
_.InterpreterDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'interpreter');
}
};
_.AffineDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'affine');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'affine.parallel') {
return this.parseParallelOp(parser, result);
}
if (op === 'affine.for') {
return this.parseForOp(parser, result);
}
if (op === 'affine.if') {
const condAttr = parser.parseAttribute();
result.addAttribute('condition', condAttr);
parser.parseDimAndSymbolList(result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
const thenRegion = result.addRegion();
const elseRegion = result.addRegion();
parser.parseRegion(thenRegion);
if (parser.parseOptionalKeyword('else')) {
parser.parseRegion(elseRegion);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'affine.apply') {
const indexType = new _.IndexType();
const mapAttr = parser.parseAttribute();
result.addAttribute('map', mapAttr);
parser.parseDimAndSymbolList(result.operands);
parser.parseOptionalAttrDict(result.attributes);
const map = mapAttr instanceof _.AffineMapAttr ? mapAttr.value : null;
const numResults = map ? map.getNumResults() : 1;
const resultTypes = [];
for (let i = 0; i < numResults; i++) {
resultTypes.push(indexType);
}
result.addTypes(resultTypes);
return true;
}
if (op === 'affine.min' || op === 'affine.max') {
const indexType = new _.IndexType();
const mapAttr = parser.parseAttribute();
result.addAttribute('map', mapAttr);
parser.parseDimAndSymbolList(result.operands);
parser.parseOptionalAttrDict(result.attributes);
result.addTypes([indexType]);
return true;
}
if (op === 'affine.store') {
return this.parseStoreOp(parser, result);
}
if (op === 'affine.load') {
return this.parseLoadOp(parser, result);
}
if (op === 'affine.vector_load') {
return this.parseVectorLoadOp(parser, result);
}
if (op === 'affine.vector_store') {
return this.parseVectorStoreOp(parser, result);
}
if (op === 'affine.prefetch') {
const indexType = new _.IndexType();
const i32Type = new _.IntegerType('i32');
const memref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseComma();
const rwSpecifier = parser.parseOptionalKeyword();
result.addAttribute('isWrite', new _.BoolAttr(rwSpecifier === 'write'));
parser.parseComma();
parser.parseKeyword('locality');
parser.parseLess();
result.addAttribute('localityHint', new _.IntegerAttr(i32Type, parser.parseInteger()));
parser.parseGreater();
parser.parseComma();
const cacheType = parser.parseOptionalKeyword();
result.addAttribute('isDataCache', new _.BoolAttr(cacheType === 'data'));
parser.parseOptionalAttrDict(result.attributes);
const type = parser.parseColonType();
parser.resolveOperand(memref, type, result.operands);
parser.resolveOperands(mapOperands, mapOperands.map(() => indexType), result.operands);
return true;
}
if (op === 'affine.dma_start') {
const indexType = new _.IndexType();
const srcMemRef = parser.parseOperand();
const srcMapOperands = [];
parser.parseAffineMapOfSSAIds(srcMapOperands, null, 'src_map', result.attributes);
parser.parseComma();
const dstMemRef = parser.parseOperand();
const dstMapOperands = [];
parser.parseAffineMapOfSSAIds(dstMapOperands, null, 'dst_map', result.attributes);
parser.parseComma();
const tagMemRef = parser.parseOperand();
const tagMapOperands = [];
parser.parseAffineMapOfSSAIds(tagMapOperands, null, 'tag_map', result.attributes);
parser.parseComma();
const numElements = parser.parseOperand();
// Parse optional stride operands
const strideOperands = parser.parseTrailingOperandList();
const types = parser.parseColonTypeList();
parser.resolveOperand(srcMemRef, types[0], result.operands);
parser.resolveOperands(srcMapOperands, srcMapOperands.map(() => indexType), result.operands);
parser.resolveOperand(dstMemRef, types[1], result.operands);
parser.resolveOperands(dstMapOperands, dstMapOperands.map(() => indexType), result.operands);
parser.resolveOperand(tagMemRef, types[2], result.operands);
parser.resolveOperands(tagMapOperands, tagMapOperands.map(() => indexType), result.operands);
parser.resolveOperand(numElements, indexType, result.operands);
parser.resolveOperands(strideOperands, strideOperands.map(() => indexType), result.operands);
return true;
}
if (op === 'affine.dma_wait') {
const indexType = new _.IndexType();
const tagMemRef = parser.parseOperand();
const tagMapOperands = [];
parser.parseAffineMapOfSSAIds(tagMapOperands, null, 'tag_map', result.attributes);
parser.parseComma();
const numElements = parser.parseOperand();
const type = parser.parseColonType();
parser.resolveOperand(tagMemRef, type, result.operands);
parser.resolveOperands(tagMapOperands, tagMapOperands.map(() => indexType), result.operands);
parser.resolveOperand(numElements, indexType, result.operands);
return true;
}
return super.parseOperation(parser, result);
}
parseForOp(parser, result) {
const indexType = new _.IndexType();
const inductionVariable = parser.parseArgument();
inductionVariable.type = indexType;
parser.parseEqual();
const numOperandsBefore = result.operands.length;
this.parseBound(true, parser, result);
const numLbOperands = result.operands.length - numOperandsBefore;
parser.parseKeyword('to');
const numOperandsBeforeUb = result.operands.length;
this.parseBound(false, parser, result);
const numUbOperands = result.operands.length - numOperandsBeforeUb;
if (parser.parseOptionalKeyword('step')) {
const step = parser.parseInteger();
result.addAttribute('step', new _.IntegerAttr(indexType, step));
} else {
result.addAttribute('step', new _.IntegerAttr(indexType, 1));
}
const regionArgs = [inductionVariable];
const operands = [];
if (parser.parseOptionalKeyword('iter_args')) {
parser.parseAssignmentList(regionArgs, operands);
result.addTypes(parser.parseArrowTypeList());
for (let i = 0; i < operands.length; i++) {
regionArgs[i + 1].type = result.types[i];
}
parser.resolveOperands(operands, result.types, result.operands);
}
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([numLbOperands, numUbOperands, operands.length]));
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseBound(isLower, parser, result) {
const indexType = new _.IndexType();
const boundAttrName = isLower ? 'lowerBoundMap' : 'upperBoundMap';
if (isLower) {
parser.parseOptionalKeyword('max');
} else {
parser.parseOptionalKeyword('min');
}
const boundOpInfos = parser.parseOperandList();
if (boundOpInfos.length > 0) {
if (boundOpInfos.length > 1) {
parser.emitError('expected only one loop bound operand');
return;
}
parser.resolveOperand(boundOpInfos[0], indexType, result.operands);
const map = new _.AffineMap(0, 1, [new _.AffineSymbolExpr(0)]);
result.addAttribute(boundAttrName, new _.AffineMapAttr(map));
return;
}
const boundAttr = parser.parseAttribute(indexType);
if (boundAttr instanceof _.AffineMapAttr) {
result.addAttribute(boundAttrName, boundAttr);
parser.parseDimAndSymbolList(result.operands);
return;
}
if (boundAttr instanceof _.IntegerAttr || typeof boundAttr === 'number' || typeof boundAttr === 'bigint') {
const value = boundAttr instanceof _.IntegerAttr ? boundAttr.value : boundAttr;
const map = new _.AffineMap(0, 0, [new _.AffineConstantExpr(Number(value))]);
result.addAttribute(boundAttrName, new _.AffineMapAttr(map));
return;
}
parser.emitError('expected valid affine map representation for loop bounds');
}
parseStoreOp(parser, result) {
const indexType = new _.IndexType();
const storeValue = parser.parseOperand();
parser.parseComma();
const memref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
const type = parser.parseColonType();
const valueType = type.elementType || type;
parser.resolveOperand(storeValue, valueType, result.operands);
parser.resolveOperand(memref, type, result.operands);
parser.resolveOperands(mapOperands, mapOperands.map(() => indexType), result.operands);
return true;
}
parseLoadOp(parser, result) {
const indexType = new _.IndexType();
const memref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
const type = parser.parseColonType();
parser.resolveOperand(memref, type, result.operands);
parser.resolveOperands(mapOperands, mapOperands.map(() => indexType), result.operands);
result.addTypes([type.elementType || type]);
return true;
}
parseVectorLoadOp(parser, result) {
const indexType = new _.IndexType();
const memref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
const memrefType = parser.parseColonType();
parser.parseComma();
const vectorType = parser.parseType();
parser.resolveOperand(memref, memrefType, result.operands);
parser.resolveOperands(mapOperands, mapOperands.map(() => indexType), result.operands);
result.addTypes([vectorType]);
return true;
}
parseVectorStoreOp(parser, result) {
const indexType = new _.IndexType();
const storeValue = parser.parseOperand();
parser.parseComma();
const memref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
const memrefType = parser.parseColonType();
parser.parseComma();
const vectorType = parser.parseType();
parser.resolveOperand(storeValue, vectorType, result.operands);
parser.resolveOperand(memref, memrefType, result.operands);
parser.resolveOperands(mapOperands, mapOperands.map(() => indexType), result.operands);
return true;
}
parseAffineMapWithMinMax(parser, result, mapOperands, mapName, groupsName, minMaxKeyword) {
parser.parseLParen();
if (parser.parseOptionalRParen()) {
result.addAttribute(mapName, new _.AffineMapAttr(new _.AffineMap(0, 0, [])));
result.addAttribute(groupsName, new _.DenseI32ArrayAttr([]));
return;
}
const flatExprs = [];
const numMapsPerGroup = [];
do {
if (parser.parseOptionalKeyword(minMaxKeyword)) {
const tmpAttrs = new Map();
parser.parseAffineMapOfSSAIds(mapOperands, null, '__map', tmpAttrs, 'Paren');
const map = tmpAttrs.get('__map');
if (!(map instanceof _.AffineMap)) {
parser.emitError(`expected affine map after '${minMaxKeyword}'`);
return;
}
for (const expr of map.getResults()) {
flatExprs.push(expr);
}
numMapsPerGroup.push(map.getNumResults());
} else {
const expr = parser.parseAffineExprOfSSAIds(mapOperands);
flatExprs.push(expr);
numMapsPerGroup.push(1);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
const flatMap = new _.AffineMap(mapOperands.length, 0, flatExprs);
result.addAttribute(mapName, new _.AffineMapAttr(flatMap));
result.addAttribute(groupsName, new _.DenseI32ArrayAttr(numMapsPerGroup));
}
parseParallelOp(parser, result) {
const ivArgs = parser.parseArgumentList('paren', false);
const indexType = new _.IndexType();
for (const iv of ivArgs) {
iv.type = indexType;
}
if (!parser.parseOptionalEqual()) {
return false;
}
const lowerMapOperands = [];
this.parseAffineMapWithMinMax(parser, result, lowerMapOperands, 'lowerBoundsMap', 'lowerBoundsGroups', 'max');
if (!parser.parseOptionalKeyword('to')) {
return false;
}
const upperMapOperands = [];
this.parseAffineMapWithMinMax(parser, result, upperMapOperands, 'upperBoundsMap', 'upperBoundsGroups', 'min');
const i64Type = new _.IntegerType('i64');
if (parser.parseOptionalKeyword('step')) {
const stepMapOperands = [];
const stepAttrs = new Map();
parser.parseAffineMapOfSSAIds(stepMapOperands, null, '__steps', stepAttrs, 'Paren');
const stepsMap = stepAttrs.get('__steps');
const steps = [];
if (stepsMap instanceof _.AffineMap) {
for (const expr of stepsMap.getResults()) {
if (expr instanceof _.AffineConstantExpr) {
steps.push(expr.value);
} else {
parser.emitError('steps must be constant integers');
return false;
}
}
}
result.addAttribute('steps', new _.ArrayAttr(steps.map((s) => new _.IntegerAttr(i64Type, s))));
} else {
const defaultSteps = [];
for (let i = 0; i < ivArgs.length; i++) {
defaultSteps.push(new _.IntegerAttr(i64Type, 1));
}
result.addAttribute('steps', new _.ArrayAttr(defaultSteps));
}
parser.resolveOperands(lowerMapOperands, lowerMapOperands.map(() => indexType), result.operands);
parser.resolveOperands(upperMapOperands, upperMapOperands.map(() => indexType), result.operands);
const reductions = [];
if (parser.parseOptionalKeyword('reduce')) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const kind = parser.parseString();
const enumValue = _.AffineDialect.symbolizeAtomicRMWKind(kind);
if (enumValue === null) {
parser.emitError(`invalid reduction value: "${kind}"`);
return false;
}
reductions.push(new _.IntegerAttr(i64Type, enumValue));
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
result.addAttribute('reductions', new _.ArrayAttr(reductions));
result.addTypes(parser.parseOptionalArrowTypeList());
const region = result.addRegion();
parser.parseRegion(region, ivArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
static symbolizeAtomicRMWKind(kind) {
const kinds = {
addf: 0, addi: 1, andi: 2, assign: 3, maximumf: 4, maxnumf: 5,
maxs: 6, maxu: 7, minimumf: 8, minnumf: 9, mins: 10, minu: 11,
mulf: 12, muli: 13, ori: 14, xori: 15
};
return Object.prototype.hasOwnProperty.call(kinds, kind) ? kinds[kind] : null;
}
};
_.MemRefDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'memref');
this.registerCustomDirective('GlobalMemrefOpTypeAndInitialValue', this.parseGlobalMemrefOpTypeAndInitialValue.bind(this));
// AtomicRMWKindAttr can appear as bare id (addi) or string ("addi") in different test files
this.registerCustomAttribute('AtomicRMWKindAttr', this.parseAtomicRMWKindAttr.bind(this));
}
parseAtomicRMWKindAttr(parser, type) {
// Accept both bare identifier (addi) and string literal ("addi")
const str = parser.parseOptionalString();
if (str !== null) {
return str;
}
if (type.values) {
const kw = parser.parseOptionalKeyword(type.values);
if (kw) {
return kw;
}
}
return null;
}
parseGlobalMemrefOpTypeAndInitialValue(parser, op, typeAttr = 'type', initialValueAttr = 'initial_value') {
const type = parser.parseType();
op.addAttribute(typeAttr, { value: type, type: 'type' });
if (parser.parseOptionalEqual()) {
if (parser.parseOptionalKeyword('uninitialized')) {
op.addAttribute(initialValueAttr, 'uninitialized');
} else {
// Pass the type to parseAttribute to suppress : type suffix parsing
const initialValue = parser.parseAttribute(type);
op.addAttribute(initialValueAttr, initialValue);
}
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'memref.tensor_load') {
const unresolvedOperands = parser.parseOperandList();
const types = parser.parseOptionalColonTypeList();
parser.resolveOperands(unresolvedOperands, types, result.operands);
// Infer result tensor type from memref operand type
if (result.operands.length > 0 && result.operands[0].type) {
const memrefType = result.operands[0].type;
// Convert memref type to tensor type
if (memrefType instanceof _.MemRefType) {
result.addTypes([new _.RankedTensorType(memrefType.shape, memrefType.elementType, null)]);
} else if (memrefType instanceof _.UnrankedMemRefType) {
result.addTypes([new _.UnrankedTensorType(memrefType.elementType)]);
}
}
return true;
}
if (op === 'memref.store') {
result.compatibility = true;
return this.parseStoreOp(parser, result);
}
if (op === 'memref.alloca_scope') {
return this.parseAllocaScopeOp(parser, result);
}
if (op === 'memref.transpose') {
return this.parseTransposeOp(parser, result);
}
if (op === 'memref.generic_atomic_rmw') {
return this.parseGenericAtomicRMWOp(parser, result);
}
if (op === 'memref.prefetch') {
const memref = parser.parseOperand();
const indices = parser.parseOperandList('square');
parser.parseComma();
const readOrWrite = parser.parseKeyword();
result.addAttribute('isWrite', readOrWrite === 'write');
parser.parseComma();
parser.parseKeyword('locality');
parser.parseLess();
result.addAttribute('localityHint', parser.parseInteger());
parser.parseGreater();
parser.parseComma();
const cacheType = parser.parseKeyword();
result.addAttribute('isDataCache', cacheType === 'data');
const type = parser.parseColonType();
parser.resolveOperand(memref, type, result.operands);
const indexType = new _.IndexType();
const indexTypes = indices.map(() => indexType);
parser.resolveOperands(indices, indexTypes, result.operands);
return true;
}
if (op === 'memref.dma_start') {
const srcMemRef = parser.parseOperand();
const srcIndices = parser.parseOperandList('square');
parser.parseComma();
const dstMemRef = parser.parseOperand();
const dstIndices = parser.parseOperandList('square');
parser.parseComma();
const numElements = parser.parseOperand();
parser.parseComma();
const tagMemRef = parser.parseOperand();
const tagIndices = parser.parseOperandList('square');
const strideInfo = [];
if (parser.parseOptionalComma()) {
let strideOp = parser.parseOptionalOperand();
while (strideOp) {
strideInfo.push(strideOp);
if (!parser.parseOptionalComma()) {
break;
}
strideOp = parser.parseOptionalOperand();
}
}
const types = parser.parseColonTypeList();
const indexType = new _.IndexType();
parser.resolveOperand(srcMemRef, types[0], result.operands);
const srcIndexTypes = srcIndices.map(() => indexType);
parser.resolveOperands(srcIndices, srcIndexTypes, result.operands);
parser.resolveOperand(dstMemRef, types[1], result.operands);
const dstIndexTypes = dstIndices.map(() => indexType);
parser.resolveOperands(dstIndices, dstIndexTypes, result.operands);
parser.resolveOperand(numElements, indexType, result.operands);
parser.resolveOperand(tagMemRef, types[2], result.operands);
const tagIndexTypes = tagIndices.map(() => indexType);
parser.resolveOperands(tagIndices, tagIndexTypes, result.operands);
if (strideInfo.length > 0) {
const strideTypes = strideInfo.map(() => indexType);
parser.resolveOperands(strideInfo, strideTypes, result.operands);
}
return true;
}
return super.parseOperation(parser, result);
}
parseGenericAtomicRMWOp(parser, result) {
const memref = parser.parseOperand();
const indices = parser.parseOperandList('square');
parser.parseColon();
const memrefType = parser.parseType();
parser.resolveOperand(memref, memrefType, result.operands);
const indexType = new _.IndexType();
const indexTypes = indices.map(() => indexType);
parser.resolveOperands(indices, indexTypes, result.operands);
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
if (memrefType && memrefType.elementType) {
result.addTypes([memrefType.elementType]);
}
return true;
}
parseTransposeOp(parser, result) {
const operand = parser.parseOperand();
const permutation = parser.parseAffineMap();
result.addAttribute('permutation', new _.AffineMapAttr(permutation));
parser.parseOptionalAttrDict(result.attributes);
const srcType = parser.parseColonType();
parser.resolveOperand(operand, srcType, result.operands);
const dstType = parser.parseKeywordType('to');
result.addTypes([dstType]);
return true;
}
parseAllocaScopeOp(parser, result) {
const resultTypes = parser.parseOptionalArrowTypeList();
result.addTypes(resultTypes);
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseStoreOp(parser, result) {
// or old: value to memref[indices] : type
let valueOperand = null;
const operand = parser.parseOptionalOperand();
if (operand) {
valueOperand = operand;
} else {
// Non-standard: constant value - store as attribute
const value = parser.parseAttribute();
result.addAttribute('value', value);
}
// Accept either ',' (new) or 'to' (old)
if (!parser.parseOptionalKeyword('to')) {
parser.parseOptionalComma();
}
const memrefOperand = parser.parseOperand();
const indices = parser.parseOperandList('square');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const memrefType = parser.parseType();
// Value type is element type of memref
const valueType = memrefType.elementType || memrefType;
if (valueOperand) {
parser.resolveOperand(valueOperand, valueType, result.operands);
}
parser.resolveOperand(memrefOperand, memrefType, result.operands);
const indexType = new _.IndexType();
parser.resolveOperands(indices, indices.map(() => indexType), result.operands);
}
return true;
}
};
_.VectorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'vector');
this.registerCustomAttribute('Vector_CombiningKindAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
this.registerCustomAttribute('Arith_FastMathAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'vector.splat') {
const unresolvedOperands = parser.parseOperandList();
const types = parser.parseOptionalColonTypeList();
parser.resolveOperands(unresolvedOperands, types, result.operands);
result.addTypes(types);
return true;
}
if (op === 'vector.contract') {
const dictAttr = parser.parseAttribute();
const lhs = parser.parseOperand();
parser.parseComma();
const rhs = parser.parseOperand();
parser.parseComma();
const acc = parser.parseOperand();
parser.parseTrailingOperandList();
parser.parseOptionalAttrDict(result.attributes);
const types = parser.parseColonTypeList();
const resultType = parser.parseKeywordType('into');
parser.resolveOperand(lhs, types[0], result.operands);
parser.resolveOperand(rhs, types[1], result.operands);
parser.resolveOperand(acc, resultType, result.operands);
result.addTypes([resultType]);
if (dictAttr instanceof _.DictionaryAttr) {
for (const [key, value] of dictAttr.value) {
result.attributes.set(key, value);
}
}
if (!result.attributes.has('kind')) {
result.addAttribute('kind', 'add');
}
return true;
}
if (op === 'vector.mask') {
let mask = null;
let passthru = null;
let hasPassthru = false;
mask = parser.parseOptionalOperand();
if (mask && parser.parseOptionalComma()) {
hasPassthru = true;
passthru = parser.parseOperand();
}
{
const region = result.addRegion();
parser.parseOptionalRegion(region);
}
parser.parseOptionalAttrDict(result.attributes);
const [maskType] = parser.parseOptionalColonTypeList();
const resultTypes = parser.parseOptionalArrowTypeList();
if (mask) {
parser.resolveOperand(mask, maskType, result.operands);
}
if (hasPassthru && passthru) {
parser.resolveOperand(passthru, resultTypes[0], result.operands);
}
result.addTypes(resultTypes);
return true;
}
if (op === 'vector.outerproduct') {
return this.parseOuterProductOp(parser, result);
}
if (op === 'vector.transfer_read' || op === 'vector.transfer_write') {
return this.parseTransferOp(parser, result);
}
if (op === 'vector.extract' && !result.isGeneric) {
result.compatibility = true;
return this.parseExtractOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseOuterProductOp(parser, result) {
const operandsInfo = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
const tLHS = parser.parseColonType();
parser.parseComma();
const tRHS = parser.parseType();
const vLHS = tLHS instanceof _.VectorType ? tLHS : null;
const vRHS = tRHS instanceof _.VectorType ? tRHS : null;
let resType = null;
if (vRHS) {
const scalableDimsRes = [
vLHS.scalableDims ? vLHS.scalableDims[0] : false,
vRHS.scalableDims ? vRHS.scalableDims[0] : false
];
resType = new _.VectorType([vLHS.shape[0], vRHS.shape[0]], vLHS.elementType, scalableDimsRes);
} else {
// Scalar RHS operand
const scalableDimsRes = [vLHS.scalableDims ? vLHS.scalableDims[0] : false];
resType = new _.VectorType([vLHS.shape[0]], vLHS.elementType, scalableDimsRes);
}
parser.resolveOperand(operandsInfo[0], tLHS, result.operands);
parser.resolveOperand(operandsInfo[1], tRHS, result.operands);
if (operandsInfo.length > 2) {
parser.resolveOperand(operandsInfo[2], resType, result.operands);
}
result.addTypes([resType]);
return true;
}
parseExtractOp(parser, result) {
// Old syntax (pre-2023): %r = vector.extract %v[0] : vector<4xf32>
// New syntax: %r = vector.extract %v[0] : f32 from vector<4xf32>
const unresolvedSource = parser.parseOperand();
const unresolvedDynIndices = [];
const indexType = new _.IndexType();
let numStaticIndices = 0;
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const staticIndex = parser.parseOptionalInteger();
if (staticIndex === null) {
const dynIndex = parser.parseOptionalOperand();
if (dynIndex) {
unresolvedDynIndices.push(dynIndex);
} else {
break;
}
} else {
numStaticIndices++;
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const resultType = parser.parseType();
if (parser.parseOptionalKeyword('from')) {
const sourceType = parser.parseType();
parser.resolveOperand(unresolvedSource, sourceType, result.operands);
result.addTypes([resultType]);
} else {
// Old syntax: the type after ':' is the source type
// Result type is inferred by removing dimensions based on indices
parser.resolveOperand(unresolvedSource, resultType, result.operands);
// Infer result type from source type and number of indices
const numIndices = numStaticIndices + unresolvedDynIndices.length;
if (resultType instanceof _.VectorType && numIndices > 0) {
const shape = resultType.shape.slice(numIndices);
if (shape.length === 0) {
// Scalar result
result.addTypes([resultType.elementType]);
} else {
const scalableDims = resultType.scalableDims ? resultType.scalableDims.slice(numIndices) : [];
result.addTypes([new _.VectorType(shape, resultType.elementType, scalableDims)]);
}
} else if (resultType instanceof _.VectorType) {
// No indices - result is the same as source
result.addTypes([resultType]);
}
}
parser.resolveOperands(unresolvedDynIndices, unresolvedDynIndices.map(() => indexType), result.operands);
}
return true;
}
parseTransferOp(parser, result) {
const op = result.name.getStringRef();
let unresolvedFirst = null;
let unresolvedSecond = null;
let unresolvedMask = null;
let indexInfo = [];
if (op === 'vector.transfer_read') {
unresolvedFirst = parser.parseOperand();
indexInfo = parser.parseOperandList('square');
parser.parseComma();
unresolvedSecond = parser.parseOperand();
if (parser.parseOptionalComma()) {
unresolvedMask = parser.parseOperand();
}
} else {
unresolvedFirst = parser.parseOperand();
parser.parseComma();
unresolvedSecond = parser.parseOperand();
indexInfo = parser.parseOperandList('square');
if (parser.parseOptionalComma()) {
unresolvedMask = parser.parseOperand();
}
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type1 = parser.parseType();
parser.parseOptionalComma();
const type2 = parser.parseType();
const indexType = new _.IndexType();
if (op === 'vector.transfer_read') {
const shapedType = type1;
const vectorType = type2;
parser.resolveOperand(unresolvedFirst, shapedType, result.operands);
parser.resolveOperands(indexInfo, indexInfo.map(() => indexType), result.operands);
const paddingType = shapedType.elementType || vectorType;
parser.resolveOperand(unresolvedSecond, paddingType, result.operands);
if (unresolvedMask) {
parser.resolveOperand(unresolvedMask, vectorType, result.operands);
}
result.addTypes([vectorType]);
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([1, indexInfo.length, 1, unresolvedMask ? 1 : 0]));
} else {
const vectorType = type1;
const shapedType = type2;
parser.resolveOperand(unresolvedFirst, vectorType, result.operands);
parser.resolveOperand(unresolvedSecond, shapedType, result.operands);
parser.resolveOperands(indexInfo, indexInfo.map(() => indexType), result.operands);
if (unresolvedMask) {
parser.resolveOperand(unresolvedMask, vectorType, result.operands);
}
if (shapedType instanceof _.RankedTensorType) {
result.addTypes([shapedType]);
}
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([1, 1, indexInfo.length, unresolvedMask ? 1 : 0]));
}
}
return true;
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'vector.shuffle') {
const maskAttr = op.attributes.get('mask');
if (maskAttr instanceof _.DenseI64ArrayAttr && op.operands.length > 0) {
const v1Type = op.operands[0].type;
if (v1Type instanceof _.VectorType) {
const maskLength = BigInt(maskAttr.value.length);
const trailingDims = v1Type.shape.slice(1);
const resultShape = [maskLength, ...trailingDims];
const resultType = new _.VectorType(resultShape, v1Type.elementType, v1Type.scalableDims ? [false, ...v1Type.scalableDims.slice(1)] : []);
op.addTypes([resultType]);
return;
}
}
}
if (op.name.getStringRef() === 'vector.to_elements' && op.operands.length > 0) {
const vecType = op.operands[0].type;
if (vecType instanceof _.VectorType) {
const elType = vecType.elementType;
for (let i = 0n; i < vecType.getNumElements(); i++) {
op.addTypes([elType]);
}
return;
}
}
super.inferResultTypes(op, vars);
}
};
_.TensorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tensor');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tensor.expand_shape') {
// The new tensor.expand_shape format includes 'output_shape':
// $src $reassociation `output_shape` custom<DynamicIndexList>(...) attr-dict `:` type($src) `into` type($result)
// Old format (deprecated):
// $src $reassociation attr-dict `:` type($src) `into` type($result)
result.compatibility = true;
const unresolvedOperand = parser.parseOperand();
const reassociation = parser.parseAttribute();
result.addAttribute('reassociation', reassociation);
if (parser.parseOptionalKeyword('output_shape')) {
// New format: parse output_shape dynamic index list
this.parseDynamicIndexList(parser, result, ['$output_shape', '$static_output_shape']);
}
// Both formats: attr-dict `:` type($src) `into` type($result)
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const srcType = parser.parseType();
parser.resolveOperands([unresolvedOperand], [srcType], result.operands);
parser.parseKeyword('into');
const resultType = parser.parseType();
result.addTypes([resultType]);
return true;
}
return super.parseOperation(parser, result);
}
};
_.torch = {};
_.torch.ListType = class extends _.Type {
constructor(containedType) {
super(null);
this.containedType = containedType;
}
static parse(parser) {
parser.parseLess();
const containedType = _.torch.TorchDialect.parseTorchDialectType(parser);
parser.parseGreater();
return new _.torch.ListType(containedType);
}
toString() {
return `!torch.list<${this.containedType}>`;
}
};
_.torch.ValueTensorType = class extends _.Type {
constructor(optionalSizes, optionalDtype, optionalSparsity) {
super(null);
this.optionalSizes = optionalSizes;
this.optionalDtype = optionalDtype;
this.optionalSparsity = optionalSparsity || null;
}
static parse(parser) {
if (!parser.parseOptionalLess()) {
return new _.torch.ValueTensorType(null, null);
}
let optionalSizes = null;
if (parser.parseOptionalStar()) {
// Unranked
} else {
parser.parseLSquare();
optionalSizes = [];
for (;;) {
if (parser.parseOptionalQuestion()) {
optionalSizes.push(-1);
} else {
const size = parser.parseOptionalInteger();
if (size === null) {
break;
}
optionalSizes.push(size);
}
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRSquare();
}
parser.parseComma();
let optionalDtype = null;
if (!parser.parseOptionalKeyword('unk')) {
optionalDtype = parser.parseType();
}
let optionalSparsity = null;
if (parser.parseOptionalComma()) {
optionalSparsity = parser.parseAttribute();
}
parser.parseGreater();
return new _.torch.ValueTensorType(optionalSizes, optionalDtype, optionalSparsity);
}
toString() {
if (this.optionalSizes === null && this.optionalDtype === null) {
return '!torch.vtensor';
}
const sizeStr = this.optionalSizes === null ? '*' :
`[${this.optionalSizes.map((s) => s === -1 ? '?' : s.toString()).join(', ')}]`;
const dtypeStr = this.optionalDtype ? this.optionalDtype.toString() : 'unk';
if (this.optionalSparsity) {
return `!torch.vtensor<${sizeStr}, ${dtypeStr}, ${this.optionalSparsity}>`;
}
return `!torch.vtensor<${sizeStr}, ${dtypeStr}>`;
}
};
_.torch.NonValueTensorType = class extends _.torch.ValueTensorType {
static parse(parser) {
if (!parser.parseOptionalLess()) {
return new _.torch.NonValueTensorType(null, null);
}
let optionalSizes = null;
if (parser.parseOptionalStar()) {
// Unranked
} else {
parser.parseLSquare();
optionalSizes = [];
for (;;) {
if (parser.parseOptionalQuestion()) {
optionalSizes.push(-1);
} else {
const size = parser.parseOptionalInteger();
if (size === null) {
break;
}
optionalSizes.push(size);
}
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRSquare();
}
parser.parseComma();
let optionalDtype = null;
if (!parser.parseOptionalKeyword('unk')) {
optionalDtype = parser.parseType();
}
let optionalSparsity = null;
if (parser.parseOptionalComma()) {
optionalSparsity = parser.parseAttribute();
}
parser.parseGreater();
return new _.torch.NonValueTensorType(optionalSizes, optionalDtype, optionalSparsity);
}
toString() {
if (this.optionalSizes === null && this.optionalDtype === null) {
return '!torch.tensor';
}
const sizeStr = this.optionalSizes === null ? '*' :
`[${this.optionalSizes.map((s) => s === -1 ? '?' : s.toString()).join(', ')}]`;
const dtypeStr = this.optionalDtype ? this.optionalDtype.toString() : 'unk';
if (this.optionalSparsity) {
return `!torch.tensor<${sizeStr}, ${dtypeStr}, ${this.optionalSparsity}>`;
}
return `!torch.tensor<${sizeStr}, ${dtypeStr}>`;
}
};
_.torch.OptionalType = class extends _.Type {
constructor(containedType) {
super(null);
this.containedType = containedType;
}
static parse(parser) {
parser.parseLess();
const containedType = _.torch.TorchDialect.parseTorchDialectType(parser);
parser.parseGreater();
return new _.torch.OptionalType(containedType);
}
toString() {
return `!torch.optional<${this.containedType}>`;
}
};
_.torch.TupleType = class extends _.Type {
constructor(containedTypes) {
super(null);
this.containedTypes = containedTypes;
}
static parse(parser) {
parser.parseLess();
const containedTypes = [];
if (!parser.parseOptionalGreater()) {
do {
containedTypes.push(_.torch.TorchDialect.parseTorchDialectType(parser));
} while (parser.parseOptionalComma());
parser.parseGreater();
}
return new _.torch.TupleType(containedTypes);
}
toString() {
return `!torch.tuple<${this.containedTypes.join(', ')}>`;
}
};
_.torch.UnionType = class extends _.Type {
constructor(containedTypes) {
super(null);
this.containedTypes = containedTypes;
}
static parse(parser) {
parser.parseLess();
const containedTypes = [];
if (!parser.parseOptionalGreater()) {
do {
containedTypes.push(_.torch.TorchDialect.parseTorchDialectType(parser));
} while (parser.parseOptionalComma());
parser.parseGreater();
}
return new _.torch.UnionType(containedTypes);
}
toString() {
return `!torch.union<${this.containedTypes.join(', ')}>`;
}
};
_.torch.DictType = class extends _.Type {
constructor(keyType, valueType) {
super(null);
this.keyType = keyType;
this.valueType = valueType;
}
static parse(parser) {
parser.parseLess();
const keyType = _.torch.TorchDialect.parseTorchDialectType(parser);
parser.parseComma();
const valueType = _.torch.TorchDialect.parseTorchDialectType(parser);
parser.parseGreater();
return new _.torch.DictType(keyType, valueType);
}
toString() {
return `!torch.dict<${this.keyType}, ${this.valueType}>`;
}
};
_.torch.NnModuleType = class extends _.Type {
constructor(className) {
super(null);
this.className = className;
}
static parse(parser) {
parser.parseLess();
const className = parser.parseString();
parser.parseGreater();
return new _.torch.NnModuleType(className);
}
toString() {
return `!torch.nn.Module<"${this.className}">`;
}
};
_.torch.TorchDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'torch');
}
static parseTorchDialectType(parser) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
switch (mnemonic) {
case 'int':
case 'float':
case 'bool':
case 'str':
case 'none':
case 'Device':
case 'Generator':
case 'qint8':
case 'quint8':
case 'qint16':
case 'qint32':
case 'quint4x2':
case 'quint2x4':
case 'LinearParams':
case 'number':
case 'any':
return new _.Type(`!torch.${mnemonic}`);
case 'list': return _.torch.ListType.parse(parser);
case 'optional': return _.torch.OptionalType.parse(parser);
case 'tuple': return _.torch.TupleType.parse(parser);
case 'union': return _.torch.UnionType.parse(parser);
case 'dict': return _.torch.DictType.parse(parser);
case 'vtensor': return _.torch.ValueTensorType.parse(parser);
case 'tensor': return _.torch.NonValueTensorType.parse(parser);
case 'nn.Module': return _.torch.NnModuleType.parse(parser);
default:
return new _.Type(`!torch.${mnemonic}`);
}
}
return parser.parseType();
}
parseType(parser) {
return _.torch.TorchDialect.parseTorchDialectType(parser);
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'torch.constant.bool' || op === 'torch.constant.device' || op === 'torch.constant.float' || op === 'torch.constant.int' || op === 'torch.constant.none' || op === 'torch.constant.number' || op === 'torch.constant.str') {
result.label = 'torch.constant';
} else if (op.startsWith('torch.aten.') || op.startsWith('torch.prim.') || op.startsWith('torch.prims.') || op.startsWith('torch.torchvision.')) {
// Collapse torch sub-dialect ops to `<subdialect>.<name>` for user display only. Parse dispatch uses the full name.
const parts = op.split('.');
result.label = `${parts[1]}.${parts[2]}`;
}
if (op === 'torch.constant.int') {
const value = parser.parseOptionalInteger();
if (value !== null) {
result.addAttribute('value', value);
}
parser.parseOptionalAttrDict(result.attributes);
result.addTypes([new _.Type('!torch.int')]);
return true;
}
if (op === 'torch.onnx.rotary_embedding') {
const unresolvedOperands = parser.parseOperandList();
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
if (parser.parseOptionalArrow()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
}
return true;
}
if (op === 'torch.bind_symbolic_shape') {
const unresolved = parser.parseOperand();
parser.parseOptionalComma();
const shapeSymbols = parser.parseOperandList('square');
parser.parseOptionalComma();
const shapeExpr = parser.parseAttribute();
result.addAttribute('shape_expressions', shapeExpr.value || shapeExpr);
parser.parseOptionalAttrDict(result.attributes);
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(unresolved, type, result.operands);
for (const sym of shapeSymbols) {
parser.resolveOperand(sym, null, result.operands);
}
return true;
}
if (op === 'torch.initialize.global_slots') {
parser.parseOptionalAttrDict(result.attributes);
parser.parseLSquare();
const slotSymNames = [];
while (!parser.parseOptionalRSquare()) {
const slotSym = parser.parseOptionalSymbolName();
slotSymNames.push(slotSym ? `@${slotSym.value || slotSym}` : '');
parser.parseLParen();
const unresolved = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(unresolved, type, result.operands);
parser.parseRParen();
}
result.addAttribute('slotSymNames', slotSymNames);
return true;
}
const opInfo = result.name.getRegisteredInfo();
if (opInfo.metadata.hasCustomAssemblyFormat && !opInfo.metadata.assemblyFormats) {
return this.parseDefaultTorchOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseDefaultTorchOp(parser, result) {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
parser.resolveOperands(unresolvedOperands, parser.parseTypeList(), result.operands);
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
if (parser.parseOptionalArrow()) {
// Handle both -> (type, type) and -> type, type syntaxes
if (parser.parseOptionalLParen()) {
const types = [];
if (!parser.parseOptionalRParen()) {
do {
types.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
result.addTypes(types);
} else {
result.addTypes(parser.parseTypeList());
}
}
{
const region = result.addRegion();
if (parser.parseOptionalRegion(region)) {
if (parser.parseOptionalKeyword('else')) {
const elseRegion = {};
parser.parseOptionalRegion(elseRegion);
result.regions.push(elseRegion);
}
}
}
return true;
}
};
_.IREEDialect = class extends _.Dialect {
constructor(operations, name) {
super(operations, name);
this.registerCustomDirective('DispatchEntryPoints', this.parseDispatchEntryPoints.bind(this));
this.registerCustomDirective('ShapedTiedResult', this.parseShapedTiedResult.bind(this));
this.registerCustomDirective('SymbolAlias', this.parseSymbolAlias.bind(this));
this.registerCustomDirective('TypeAlias', this.parseTypeAlias.bind(this));
this.registerCustomDirective('WorkgroupCountRegion', this.parseWorkgroupCountRegion.bind(this));
this.registerCustomDirective('ShapedFunctionType', this.parseShapedFunctionType.bind(this));
}
parseShapedFunctionType(parser, op, unresolvedArguments /*, otherArgs */) {
const operandTypes = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const type = parser.parseType();
if (type) {
operandTypes.push(type);
}
if (parser.parseOptionalLBrace()) {
// Skip balanced braces content (dynamic dims)
let depth = 1;
while (depth > 0) {
if (parser.parseOptionalLBrace()) {
depth++;
} else if (parser.parseOptionalRBrace()) {
depth--;
} else {
parser.parseOperand();
parser.parseOptionalComma();
}
}
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseArrow();
const operands = unresolvedArguments || op.operands;
const resultTypes = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
parser.parseShapedResultList(operands, operandTypes, resultTypes, null, op.operands);
parser.parseRParen();
}
} else {
parser.parseShapedResultList(operands, operandTypes, resultTypes, null, op.operands);
}
op.addTypes(resultTypes);
}
parseDispatchEntryPoints(parser, op, attrName = 'entry_points') {
const entryPoints = [];
if (parser.parseOptionalLBrace()) {
do {
const sym = parser.parseOptionalSymbolName();
if (sym) {
let symbol = `@${sym.value || sym}`;
if (parser.parseOptionalColon()) {
if (parser.parseOptionalColon()) {
const nestedSym = parser.parseOptionalSymbolName();
if (nestedSym) {
symbol += `::@${nestedSym.value || nestedSym}`;
}
}
}
entryPoints.push(symbol);
}
} while (parser.parseOptionalComma());
parser.parseRBrace();
} else {
const sym = parser.parseOptionalSymbolName();
if (sym) {
let symbol = `@${sym.value || sym}`;
if (parser.parseOptionalColon()) {
if (parser.parseOptionalColon()) {
const nestedSym = parser.parseOptionalSymbolName();
if (nestedSym) {
symbol += `::@${nestedSym.value || nestedSym}`;
}
}
}
entryPoints.push(symbol);
}
}
const value = entryPoints.length === 1 ? entryPoints[0] : entryPoints;
op.addAttribute(attrName, value);
}
parseShapedTiedResult(parser, op /*, args */) {
// or just: type{dims}
const tiedOp = parser.parseOptionalOperand();
if (tiedOp) {
parser.parseKeyword('as');
}
const resultType = parser.parseType();
op.types.push(resultType); // Only add the type
if (parser.parseOptionalLBrace()) {
const indexType = new _.IndexType();
if (!parser.parseOptionalRBrace()) {
do {
const dim = parser.parseOptionalOperand();
if (dim) {
parser.resolveOperand(dim, indexType, op.operands);
} else {
break;
}
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
}
parseSymbolAlias(parser, op, symNameAttr, aliasAttr) {
const aliasSym = parser.parseOptionalSymbolName();
const alias = aliasSym ? `@${aliasSym.value || aliasSym}` : '';
let symName = alias;
if (parser.parseOptionalKeyword('as')) {
if (parser.parseOptionalLParen()) {
const str = parser.parseOptionalString();
if (str === null) {
const sym = parser.parseOptionalSymbolName();
if (sym) {
symName = `@${sym.value || sym}`;
}
} else {
symName = str;
}
parser.parseOptionalRParen();
}
}
if (symNameAttr && aliasAttr) {
op.addAttribute(symNameAttr, symName);
op.addAttribute(aliasAttr, alias);
}
}
parseTypeAlias(parser, op, encodingAttrName, typeArg) {
const encodingType = parser.parseType();
let storageType = encodingType;
if (parser.parseOptionalKeyword('as')) {
storageType = parser.parseType();
}
if (encodingAttrName) {
op.addAttribute(encodingAttrName, encodingType);
}
if (!Array.isArray(typeArg)) {
throw new mlir.Error(`Invalid argument 'typeArg'.`);
}
if (typeArg.length > 0) {
typeArg[0] = storageType;
} else {
typeArg.push(storageType);
}
}
parseWorkgroupCountRegion(parser, op /*, workgroupCountName */) {
if (!parser.parseOptionalKeyword('workgroups')) {
return;
}
const region = { blocks: [] };
const block = { arguments: [], operations: [] };
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
if (parser.parseOptionalColon()) {
arg.type = parser.parseType();
}
block.arguments.push(arg);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalArrow()) {
parser.parseCommaSeparatedList('paren', () => {
parser.parseType();
});
}
if (parser.parseOptionalLBrace()) {
while (!parser.parseOptionalRBrace()) {
const innerOp = parser.parseOperation();
if (innerOp) {
block.operations.push(innerOp);
}
}
}
region.blocks.push(block);
op.regions.push(region);
}
};
_.iree = {};
_.iree.DispatchTensorType = class extends _.Type {
constructor(access, boundType) {
super(null);
this.access = access;
this.boundType = boundType;
}
static parse(parser) {
parser.parseLess();
const access = parser.parseKeyword();
parser.parseColon();
const boundType = parser.parseOptionalType();
if (boundType) {
parser.parseGreater();
return new _.iree.DispatchTensorType(access, boundType);
}
const dims = parser.parseDimensionList().dimensions;
const elementType = parser.parseType();
parser.parseGreater();
return new _.iree.DispatchTensorType(access, new _.RankedTensorType(dims, elementType));
}
toString() {
return `!flow.dispatch.tensor<${this.access}:${this.boundType}>`;
}
};
_.HALDialect = class extends _.IREEDialect {
constructor(operations) {
super(operations, 'hal');
this.simpleTypes = new Set(['allocator', 'buffer', 'buffer_view', 'channel', 'command_buffer', 'descriptor_set', 'descriptor_set_layout', 'device', 'event', 'executable', 'executable_layout', 'fence', 'file', 'semaphore']);
this.registerCustomAttribute('HAL_PipelineLayoutAttr', this.parsePipelineLayoutAttr.bind(this));
this.registerCustomDirective('ExportConditionRegion', this.parseExportConditionRegion.bind(this));
this.registerCustomDirective('TargetConditionObjects', this.parseTargetConditionObjects.bind(this));
this.registerCustomDirective('WorkgroupCountRegion', this.parseWorkgroupCountRegion.bind(this));
this.registerCustomDirective('DeviceQueueAffinityList', this.parseDeviceQueueAffinityList.bind(this));
this.registerCustomDirective('Bindings', this.parseBindingList.bind(this));
this.registerCustomDirective('BindingTable', this.parseBindingList.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (this.simpleTypes.has(mnemonic)) {
// Note: !hal.buffer{%size} syntax is handled by custom<SizeAwareType> directive,
// not by the type parser. The type parser just returns the base type.
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'hal.tensor.cast') {
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
const type = parser.parseType();
parser.resolveOperands(unresolvedOperands, [type], result.operands);
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
if (op === 'hal.constant') {
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const value = parser.parseAttribute();
result.addAttribute('value', value.value === undefined ? value : value.value);
result.addTypes(parser.parseOptionalColonTypeList());
return true;
}
if (op === 'hal.device.switch') {
if (parser.parseOptionalLess()) {
while (!parser.parseOptionalGreater()) {
const operand = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
parser.parseOptionalComma();
}
}
if (parser.parseOptionalArrow() || parser.parseOptionalColon()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
}
for (;;) {
const caseAttr = parser.parseOptionalAttribute();
if (!caseAttr) {
break;
}
const region = {};
region.caseAttribute = caseAttr;
const regionObj = {};
if (parser.parseOptionalRegion(regionObj)) {
Object.assign(region, regionObj);
}
result.regions.push(region);
parser.parseOptionalComma();
}
return true;
}
if (op === 'hal.executable.constant.block') {
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOptionalOperand();
if (arg) {
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(arg, type, result.operands);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalArrow()) {
const resultTypes = [];
const resultAttrs = [];
parser.parseFunctionResultList(resultTypes, resultAttrs);
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType([], resultTypes)));
}
if (parser.parseOptionalKeyword('as')) {
if (parser.parseOptionalLParen()) {
const keys = [];
if (!parser.parseOptionalRParen()) {
do {
const str = parser.parseOptionalString();
if (str !== null) {
keys.push(str);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
result.addAttribute('keys', keys);
} else {
const key = parser.parseOptionalString();
if (key !== null) {
result.addAttribute('keys', [key]);
}
}
}
{
const region = result.addRegion();
parser.parseOptionalRegion(region);
}
return true;
}
// Handle hal.executable.create with both old (layouts) and new (affinity) syntax
if (op === 'hal.executable.create') {
for (;;) {
const paramName = parser.parseOptionalKeyword();
if (!paramName) {
break;
}
const normalizedName = paramName === 'layouts' ? 'affinity' : paramName;
this.parseParamValue(parser, result,normalizedName);
}
result.addTypes(parser.parseOptionalColonTypeList());
return true;
}
if (op === 'hal.allocator.allocate') {
this.parseAngleBracketOperand(parser, result);
this.parseNamedAttr(parser, result,'affinity');
this.parseNamedAttr(parser, result,'type');
this.parseNamedAttr(parser, result,'usage');
result.addTypes(parser.parseOptionalColonTypeList());
if (parser.parseOptionalLBrace()) {
parser.parseOptionalOperand();
parser.parseRBrace();
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.allocator.compute_size') {
this.parseAngleBracketOperand(parser, result);
this.parseNamedAttr(parser, result,'shape');
this.parseNamedAttr(parser, result,'type');
this.parseNamedAttr(parser, result,'encoding');
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.buffer_view.create') {
if (parser.parseOptionalKeyword('buffer')) {
parser.parseLParen();
const operand = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.parseRParen();
parser.resolveOperand(operand, type, result.operands);
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const op = parser.parseOptionalOperand();
if (op) {
parser.resolveOperand(op, null, result.operands);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
}
this.parseNamedAttr(parser, result,'shape');
this.parseNamedAttr(parser, result,'type');
this.parseNamedAttr(parser, result,'encoding');
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.command_buffer.begin' || op === 'hal.command_buffer.end' || op === 'hal.command_buffer.device') {
this.parseAngleBracketOperand(parser, result);
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.command_buffer.create') {
this.parseNamedAttr(parser, result,'device');
this.parseNamedAttr(parser, result,'mode');
this.parseNamedAttr(parser, result,'categories');
this.parseNamedAttr(parser, result,'affinity');
this.parseNamedAttr(parser, result,'bindings');
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.command_buffer.dispatch') {
this.parseAngleBracketOperand(parser, result);
if (parser.parseOptionalKeyword('target')) {
parser.parseLParen();
const operand = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.parseRParen();
parser.resolveOperand(operand, type, result.operands);
if (parser.parseOptionalLSquare()) {
parser.parseOptionalOperand();
parser.parseOptionalAttribute();
parser.parseRSquare();
}
}
this.parseNamedAttr(parser, result,'workgroups');
this.parseNamedAttr(parser, result,'constants');
if (parser.parseOptionalKeyword('bindings')) {
parser.parseLParen();
this.parseBindings(parser, result);
parser.parseRParen();
}
this.parseNamedAttr(parser, result,'flags');
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.command_buffer.dispatch.symbol') {
this.parseAngleBracketOperand(parser, result);
this.parseNamedAttr(parser, result,'target');
this.parseNamedAttr(parser, result,'workgroups');
this.parseNamedAttr(parser, result,'constants');
if (parser.parseOptionalKeyword('bindings')) {
parser.parseLParen();
this.parseBindings(parser, result);
parser.parseRParen();
}
this.parseNamedAttr(parser, result,'flags');
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.command_buffer.push_descriptor_set') {
this.parseAngleBracketOperand(parser, result);
if (parser.parseOptionalKeyword('layout')) {
parser.parseLParen();
const operand = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.parseRParen();
parser.resolveOperand(operand, type, result.operands);
if (parser.parseOptionalLSquare()) {
parser.parseOptionalOperand();
parser.parseRSquare();
}
}
if (parser.parseOptionalKeyword('bindings')) {
parser.parseLParen();
this.parseBindings(parser, result);
parser.parseRParen();
}
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.descriptor_set_layout.create' || op === 'hal.executable_layout.create' || op === 'hal.executable_layout.lookup') {
for (;;) {
const paramName = parser.parseOptionalKeyword();
if (!paramName) {
break;
}
this.parseParamValue(parser, result,paramName);
}
result.addTypes(parser.parseOptionalColonTypeList());
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.device.query') {
this.parseAngleBracketOperand(parser, result);
this.parseNamedAttr(parser, result,'key');
result.addTypes(parser.parseOptionalColonTypeList());
if (parser.parseOptionalEqual()) {
const value = parser.parseAttribute();
result.addAttribute('default', value.value);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'hal.executable.condition' || op === 'hal.executable.constant.block') {
const sig = parser.parseFunctionSignatureWithArguments(false);
const argTypes = sig.arguments.map((a) => a.type);
const type = new _.FunctionType(argTypes, sig.resultTypes);
result.addAttribute('function_type', new _.TypeAttrOf(type));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
{
const region = result.addRegion();
parser.parseOptionalRegion(region, sig.arguments);
}
return true;
}
// Handle operations with visibility + symbol (similar to flow dialect)
if (op === 'hal.executable' || op === 'hal.executable.source' || op === 'hal.interface' || op === 'hal.executable.binary') {
this.parseSymbolVisibility(parser, result);
const sym = parser.parseOptionalSymbolName();
if (sym) {
result.addAttribute('sym_name', sym.value || sym);
}
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
{
const region = result.addRegion();
parser.parseOptionalRegion(region);
}
return true;
}
const subspanSym = op === 'hal.interface.binding.subspan' ? parser.parseOptionalSymbolName() : null;
if (subspanSym) {
// Old syntax - parse symbol reference and bracket expression
result.addAttribute('layout', `@${subspanSym.value || subspanSym}`);
const unresolvedOperands = [];
const indexType = new _.IndexType();
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedOperands.push(operand);
} else {
// Skip non-operand tokens (like ::@nested)
parser.parseAttribute();
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => indexType), result.operands);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type) {
result.addTypes([type]);
}
if (parser.parseOptionalLBrace()) {
const dynamicDimOperands = [];
if (!parser.parseOptionalRBrace()) {
do {
const dimOperand = parser.parseOperand();
dynamicDimOperands.push(dimOperand);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
parser.resolveOperands(dynamicDimOperands, dynamicDimOperands.map(() => indexType), result.operands);
}
}
return true;
}
// Old IREE hal.interface.binding: visibility @sym, key=value, key="string", ...
if (op === 'hal.interface.binding') {
this.parseSymbolVisibility(parser, result);
const sym = parser.parseOptionalSymbolName();
if (sym) {
result.addAttribute('sym_name', sym.value || sym);
}
while (parser.parseOptionalComma()) {
const key = parser.parseKeyword();
parser.parseEqual();
const value = parser.parseOptionalAttribute();
if (value === null) {
const str = parser.parseOptionalString();
result.addAttribute(key, str === null ? parser.parseKeyword() : str);
} else {
result.addAttribute(key, value.value === undefined ? value : value.value);
}
}
return true;
}
// hal.executable.variant:
// New: @sym target(#attr) objects(...)? sources(...)? attr-dict-with-keyword { body }
// Old: @sym, target = #attr attributes {...} { body }
if (op === 'hal.executable.variant') {
this.parseSymbolVisibility(parser, result);
const sym = parser.parseOptionalSymbolName();
if (sym) {
result.addAttribute('sym_name', sym.value || sym);
}
if (parser.parseOptionalComma()) {
// Old syntax: , target = #attr
const key = parser.parseKeyword();
parser.parseEqual();
const value = parser.parseAttribute();
result.addAttribute(key, value.value === undefined ? value : value.value);
} else if (parser.parseOptionalKeyword('target')) {
// New syntax: target(#attr)
parser.parseLParen();
const value = parser.parseAttribute();
result.addAttribute('target', value.value === undefined ? value : value.value);
parser.parseRParen();
if (parser.parseOptionalKeyword('objects')) {
parser.parseLParen();
const objects = parser.parseAttribute();
result.addAttribute('objects', objects.value === undefined ? objects : objects.value);
parser.parseRParen();
}
if (parser.parseOptionalKeyword('sources')) {
parser.parseLParen();
const sources = parser.parseAttribute();
result.addAttribute('sources', sources.value === undefined ? sources : sources.value);
parser.parseRParen();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseOptionalRegion(region);
return true;
}
// Old IREE hal.executable.entry_point: visibility @sym attributes {...} { optional region }
if (op === 'hal.executable.entry_point') {
this.parseSymbolVisibility(parser, result);
const sym = parser.parseOptionalSymbolName();
if (sym) {
result.addAttribute('sym_name', sym.value || sym);
}
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
const region = result.addRegion();
parser.parseOptionalRegion(region);
return true;
}
return super.parseOperation(parser, result);
}
parsePipelineLayoutAttr(parser) {
// HAL_PipelineLayoutAttr format: <constants = N, bindings = [...], flags = ...>
// Try parseAttribute first (handles <...> syntax), fall back to optional
const attr = parser.parseOptionalAttribute();
return attr;
}
// Parses [(%op : type)[off, sz], ...] or [%idx = (%op : type)[off, sz], ...]
parseBindings(parser, result) {
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
if (parser.parseOptionalOperand()) {
parser.parseEqual();
}
parser.parseLParen();
const bindOp = parser.parseOperand();
let bindType = null;
if (parser.parseOptionalColon()) {
bindType = parser.parseType();
}
parser.parseRParen();
parser.resolveOperand(bindOp, bindType, result.operands);
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const sizeOp = parser.parseOperand();
parser.resolveOperand(sizeOp, null, result.operands);
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
// Parses <%operand : type> from old IREE HAL format.
parseAngleBracketOperand(parser, result) {
if (parser.parseOptionalLess()) {
const operand = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
parser.parseGreater();
}
}
parseNamedAttr(parser, result, name) {
if (parser.parseOptionalKeyword(name)) {
this.parseParamValue(parser, result, name);
}
}
parseParamValue(parser, result, name) {
parser.parseLParen();
if (name === 'key') {
const category = parser.parseString();
result.addAttribute('category', category);
parser.parseColon();
parser.parseColon();
const key = parser.parseString();
result.addAttribute(name, key);
} else {
const operand = parser.parseOptionalOperand();
if (operand) {
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
while (parser.parseOptionalComma()) {
const nextOp = parser.parseOperand();
let nextType = null;
if (parser.parseOptionalColon()) {
nextType = parser.parseType();
}
parser.parseOptionalEllipsis();
parser.resolveOperand(nextOp, nextType, result.operands);
}
} else if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
const firstOp = parser.parseOptionalOperand();
if (firstOp) {
const ops = [firstOp];
while (parser.parseOptionalComma()) {
ops.push(parser.parseOperand());
}
const indexType = new _.IndexType();
parser.resolveOperands(ops, ops.map(() => indexType), result.operands);
} else {
const value = parser.parseAttribute();
result.addAttribute(name, value);
while (parser.parseOptionalComma()) {
parser.parseAttribute();
}
}
parser.parseRSquare();
}
} else {
const value = parser.parseAttribute();
result.addAttribute(name, value);
}
}
parser.parseRParen();
}
parseExportConditionRegion(parser, op /*, conditionName */) {
parser.parseLParen();
const regionArgs = [];
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
regionArgs.push({ value: arg, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseArrow();
parser.parseType();
const region = { arguments: regionArgs };
parser.parseRegion(region);
op.regions.push(region);
}
parseTargetConditionObjects(parser, op /*, targetsAttr, targetOrdinalsAttr, targetObjectsAttr, targetRegionsName */) {
// #target if(...) { region } ordinal(N) = [objects], ...
do {
const attr = parser.parseOptionalAttribute();
if (attr) {
// hash_identifier attribute consumed
}
if (parser.parseOptionalKeyword('if')) {
this.parseTargetConditionRegion(parser, op);
}
if (parser.parseOptionalKeyword('ordinal')) {
parser.parseCommaSeparatedList('paren', () => {
parser.parseInteger();
});
}
if (parser.parseOptionalEqual()) {
if (parser.parseOptionalLSquare()) {
// Skip balanced bracket content
let depth = 1;
while (depth > 0) {
if (parser.parseOptionalLSquare()) {
depth++;
} else if (parser.parseOptionalRSquare()) {
depth--;
} else {
parser.parseAttribute();
parser.parseOptionalComma();
}
}
}
}
} while (parser.parseOptionalComma());
}
parseTargetConditionRegion(parser, op) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
parser.parseOperand();
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
if (parser.parseOptionalArrow()) {
parser.parseType();
}
{
const region = op.addRegion();
parser.parseOptionalRegion(region);
}
}
parseWorkgroupCountRegion(parser, op /*, workgroupCountName */) {
// (args) -> (index, index, index) { region }
const region = { blocks: [] };
const block = { arguments: [], operations: [] };
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
if (parser.parseOptionalColon()) {
arg.type = parser.parseType();
}
block.arguments.push(arg);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalArrow()) {
parser.parseCommaSeparatedList('paren', () => {
parser.parseType();
});
}
region.blocks.push(block);
const regionObj = {};
if (parser.parseOptionalRegion(regionObj)) {
Object.assign(region, regionObj);
}
op.regions.push(region);
}
parseDeviceQueueAffinityList(parser, op, devicesOperands, devicesTypes, queueAffinitiesOperands) {
parser.parseLSquare();
do {
parser.parseLParen();
devicesOperands.push(parser.parseOperand());
parser.parseComma();
queueAffinitiesOperands.push(parser.parseOperand());
parser.parseColon();
devicesTypes.push(parser.parseType());
parser.parseComma();
parser.parseType();
parser.parseRParen();
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parseBindingList(parser, op, bindingBuffersOperands, bindingBuffersTypes, bindingOffsetsOperands, bindingLengthsOperands) {
do {
parser.parseLParen();
bindingBuffersOperands.push(parser.parseOperand());
parser.parseColon();
bindingBuffersTypes.push(parser.parseType());
parser.parseRParen();
parser.parseLSquare();
bindingOffsetsOperands.push(parser.parseOperand());
parser.parseComma();
bindingLengthsOperands.push(parser.parseOperand());
parser.parseRSquare();
} while (parser.parseOptionalComma());
}
};
_.IREECodegenDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'iree_codegen');
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'iree_codegen.inner_tiled') {
const outputsEntry = vars.get('outputs');
if (outputsEntry.types.length > 0) {
op.addTypes(outputsEntry.types);
return;
}
}
super.inferResultTypes(op, vars);
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'iree_codegen.workgroup_count_hint') {
const staticSizes = [];
const unresolvedSizes = [];
parser.parseOptionalKeyword('sizes');
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedSizes.push(operand);
staticSizes.push(-9223372036854775808n);
} else {
staticSizes.push(parser.parseInteger('int64'));
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
// Resolve operands with index type (sizes are typically index)
const indexType = new _.IndexType();
for (const unresolved of unresolvedSizes) {
parser.resolveOperand(unresolved, indexType, result.operands);
}
if (staticSizes.length > 0) {
result.addAttribute('static_sizes', staticSizes);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
return super.parseOperation(parser, result);
}
};
_.HALLoaderDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'hal_loader');
this.registerCustomDirective('DispatchBindings', this.parseDispatchBindings.bind(this));
}
parseDispatchBindings(parser, op, bindingBuffersOperands, bindingBuffersTypes, bindingOffsetsOperands, bindingLengthsOperands) {
do {
parser.parseLParen();
bindingBuffersOperands.push(parser.parseOperand());
parser.parseColon();
bindingBuffersTypes.push(parser.parseType());
parser.parseRParen();
parser.parseLSquare();
bindingOffsetsOperands.push(parser.parseOperand());
parser.parseComma();
bindingLengthsOperands.push(parser.parseOperand());
parser.parseRSquare();
} while (parser.parseOptionalComma());
}
};
_.UtilDialect = class extends _.IREEDialect {
constructor(operations) {
super(operations, 'util');
this.registerCustomDirective('OperandTypeList', this.parseOperandTypeList.bind(this));
this.registerCustomDirective('TiedFunctionResultList', this.parseTiedFunctionResultList.bind(this));
this.registerCustomDirective('TypeAlias', this.parseTypeAlias.bind(this));
this.registerCustomDirective('TypedValueList', this.parseTypedValueList.bind(this));
this.registerCustomDirective('RangeList', this.parseRangeList.bind(this));
this.registerCustomDirective('ListTypeGet', this.parseListTypeGet.bind(this));
this.registerCustomDirective('ListTypeSet', this.parseListTypeSet.bind(this));
this.registerCustomDirective('ValueTypeList', this.parseValueTypeList.bind(this));
this.simpleTypes = new Set(['buffer', 'list', 'object', 'ptr']);
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'list') {
return _.util.ListType.parse(parser);
}
if (mnemonic === 'ptr') {
parser.parseLess();
const targetType = parser.parseType();
parser.parseGreater();
return new _.Type(`!${dialect}.ptr<${targetType}>`);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseTypeAlias(parser /*, op, args */) {
parser.parseType();
if (parser.parseOptionalKeyword('as')) {
parser.parseType();
}
}
parseTypedValueList(parser, op /*, args */) {
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
const unresolvedValues = [];
do {
unresolvedValues.push(parser.parseOperand());
} while (parser.parseOptionalComma());
for (const unresolved of unresolvedValues) {
parser.resolveOperand(unresolved, null, op.operands);
}
parser.parseRSquare();
}
}
parseOperandTypeList(parser, op /*, args */) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
let index = 0;
do {
const type = parser.parseType();
if (index < op.operands.length) {
op.operands[index].type = type;
}
index++;
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parseTiedFunctionResultList(parser, op /*, args */) {
const parseTiedResultOrType = () => {
const tiedRef = parser.parseOptionalOperand();
if (tiedRef) {
let tiedType = null;
for (let i = 0; i < op.operands.length; i++) {
if (op.operands[i].value === tiedRef) {
tiedType = op.operands[i].type;
break;
}
}
if (parser.parseOptionalKeyword('as')) {
return parser.parseType();
}
if (tiedType) {
return tiedType;
}
return new _.Type('!util.unknown');
}
return parser.parseType();
};
if (parser.parseOptionalLParen()) {
let index = 0;
if (!parser.parseOptionalRParen()) {
do {
const type = parseTiedResultOrType();
if (index < op.types.length) {
op.types[index] = type;
} else {
op.addTypes([type]);
}
index++;
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
let index = 0;
do {
const type = parseTiedResultOrType();
if (index < op.types.length) {
op.types[index] = type;
} else {
op.addTypes([type]);
}
index++;
} while (parser.parseOptionalComma());
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'util.assume.int') {
return this.parseAssumeIntOp(parser, result);
}
if (op === 'util.initializer') {
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'util.func') {
this.parseUtilFuncOp(parser, result);
return true;
}
if (op === 'util.hoistable_conversion') {
return this.parseHoistableConversionOp(parser, result);
}
if (op === 'util.unfoldable_constant') {
parser.parseOptionalAttrDict(result.attributes);
const value = parser.parseAttribute();
result.addAttribute('value', value);
// Use type from attribute if present, otherwise parse explicit type
if (value && value.type) {
result.addTypes([value.type]);
} else if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
}
return true;
}
return super.parseOperation(parser, result);
}
parseHoistableConversionOp(parser, result) {
const tag = parser.parseString();
result.addAttribute('tag', tag);
parser.parseKeyword('inverts');
parser.parseLParen();
const inverseTag = parser.parseString();
result.addAttribute('inverseTag', inverseTag);
parser.parseRParen();
const regionArgs = [];
const inputOperands = [];
parser.parseAssignmentList(regionArgs, inputOperands);
const fnType = parser.parseColonType();
for (let i = 0; i < regionArgs.length; i++) {
regionArgs[i].type = fnType.inputs[i];
}
parser.resolveOperands(inputOperands, fnType.inputs, result.operands);
result.addTypes(fnType.results);
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseUtilFuncOp(parser, result) {
parser.parseOptionalVisibilityKeyword(result.attributes);
parser.parseSymbolName('sym_name', result.attributes);
const argResult = parser.parseFunctionArgumentList(false);
const resultTypes = [];
const resultAttrs = [];
const tiedOperandIndices = [];
// Parse result list which may contain:
// - Regular type: tensor<...>
// - Tied reference: %arg1 (inherits type from argument)
// - Tied with type override: %arg2 as tensor<...>
const parseTiedResultOrType = () => {
const tiedRef = parser.parseOptionalOperand();
if (tiedRef) {
let tiedIndex = -1;
for (let i = 0; i < argResult.arguments.length; i++) {
if (argResult.arguments[i].value === tiedRef) {
tiedIndex = i;
break;
}
}
tiedOperandIndices.push(tiedIndex);
// Check for 'as type' override
if (parser.parseOptionalKeyword('as')) {
return parser.parseType();
}
if (tiedIndex >= 0 && argResult.arguments[tiedIndex].type) {
return argResult.arguments[tiedIndex].type;
}
return new _.Type('!util.unknown');
}
tiedOperandIndices.push(-1);
return parser.parseType();
};
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
resultTypes.push(parseTiedResultOrType());
const attrList = new Map();
parser.parseOptionalAttrDict(attrList);
resultAttrs.push(attrList.size > 0 ? attrList : null);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
do {
resultTypes.push(parseTiedResultOrType());
resultAttrs.push(null);
} while (parser.parseOptionalComma());
}
}
const argTypes = argResult.arguments.filter((a) => a.value !== '...').map((a) => a.type);
const type = new _.FunctionType(argTypes, resultTypes);
result.addAttribute('function_type', new _.TypeAttrOf(type));
if (tiedOperandIndices.some((i) => i >= 0)) {
result.addAttribute('tied_operands', tiedOperandIndices);
}
if (resultAttrs.some((a) => a !== null)) {
result.addAttribute('res_attrs', resultAttrs);
}
const argAttrs = argResult.arguments.filter((a) => a.value !== '...').map((a) => a.attrs || null);
if (argAttrs.some((a) => a !== null)) {
result.addAttribute('arg_attrs', argAttrs);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion(), argResult.arguments);
}
parseAssumeIntOp(parser, result) {
const allOperandAssumptions = [];
const unresolvedOperands = [];
do {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
const operandAssumptions = [];
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const assumption = this.parseIntAssumptionAttr(parser);
operandAssumptions.push(assumption);
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
} else if (parser.parseOptionalLess()) {
// parseIntAssumptionAttr will call parseLess() but we already consumed it
// Re-structure: parse content between < and >
const assumption = {};
if (!parser.parseOptionalGreater()) {
do {
const key = parser.parseKeyword();
parser.parseEqual();
assumption[key] = parser.parseInteger();
} while (parser.parseOptionalComma());
parser.parseGreater();
}
operandAssumptions.push(assumption);
}
allOperandAssumptions.push(operandAssumptions);
} while (parser.parseOptionalComma());
parser.parseColon();
const parsedOperandTypes = [];
do {
const type = parser.parseType();
parsedOperandTypes.push(type);
} while (parser.parseOptionalComma());
parser.resolveOperands(unresolvedOperands, parsedOperandTypes, result.operands);
for (const type of parsedOperandTypes) {
result.addTypes([type || null]);
}
result.addAttribute('assumptions', allOperandAssumptions);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseIntAssumptionAttr(parser) {
parser.parseLess();
const assumption = {};
if (!parser.parseOptionalGreater()) {
do {
const key = parser.parseKeyword();
if (!parser.parseOptionalEqual()) {
parser.emitError(`Expected '=' after ${key}`);
}
assumption[key] = parser.parseInteger();
} while (parser.parseOptionalComma());
parser.parseGreater();
}
return assumption;
}
parseRangeList(parser, op, offsetsAttr) {
const unresolvedOffsets = [];
const unresolvedLengths = [];
do {
parser.parseLSquare();
unresolvedOffsets.push(parser.parseOperand());
parser.parseKeyword('for');
unresolvedLengths.push(parser.parseOperand());
parser.parseRSquare();
} while (parser.parseOptionalComma());
const indexType = new _.IndexType();
for (const unresolved of unresolvedOffsets) {
parser.resolveOperand(unresolved, indexType, op.operands);
}
for (const unresolved of unresolvedLengths) {
parser.resolveOperand(unresolved, indexType, op.operands);
}
if (offsetsAttr) {
op.addAttribute(`${offsetsAttr}_count`, unresolvedOffsets.length);
}
}
parseListTypeGet(parser, op, listTypeArr, resultTypeArr) {
const listType = parser.parseType();
let elementType = null;
if (parser.parseOptionalArrow()) {
elementType = parser.parseType();
} else if (listType instanceof _.util.ListType) {
elementType = listType.elementType;
}
if (Array.isArray(listTypeArr) && listType) {
listTypeArr.push(listType);
}
if (Array.isArray(resultTypeArr) && elementType) {
resultTypeArr.push(elementType);
}
}
parseListTypeSet(parser, op, listTypeArr, valueTypeArr) {
const leadingType = parser.parseType();
let listType = null;
let elementType = null;
if (parser.parseOptionalArrow()) {
elementType = leadingType;
listType = parser.parseType();
} else if (leadingType instanceof _.util.ListType) {
listType = leadingType;
elementType = leadingType.elementType;
}
if (Array.isArray(listTypeArr) && listType) {
listTypeArr.push(listType);
}
if (Array.isArray(valueTypeArr) && elementType) {
valueTypeArr.push(elementType);
}
}
parseValueTypeList(parser, op, valuesOperands, valuesTypes) {
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
valuesOperands.push(parser.parseOperand());
parser.parseColon();
valuesTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
};
_.FlowDialect = class extends _.IREEDialect {
constructor(operations) {
super(operations, 'flow');
this.registerCustomDirective('DispatchWorkgroupBody', this.parseDispatchWorkgroupBody.bind(this));
this.registerCustomDirective('DispatchWorkgroupsCountRegion', this.parseDispatchWorkgroupsCountRegion.bind(this));
this.registerCustomDirective('ShapedFunctionType', this.parseShapedFunctionType.bind(this));
this.registerCustomDirective('ShapedOperandList', this.parseShapedOperandList.bind(this));
this.registerCustomDirective('ParameterReference', this.parseParameterReference.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'channel') {
return new _.Type(`!${dialect}.${mnemonic}`);
}
if (mnemonic === 'dispatch.tensor') {
return _.iree.DispatchTensorType.parse(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseParameterReference(parser, op, scopeOperands, keyOperands) {
const firstOperand = parser.parseOperand();
if (parser.parseOptionalColon()) {
parser.parseColon();
const keyOperand = parser.parseOperand();
scopeOperands.push(firstOperand);
keyOperands.push(keyOperand);
} else {
keyOperands.push(firstOperand);
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'flow.ex.stream.fragment') {
return this.parseDispatchWorkgroupsOp(parser, result);
}
if (op === 'flow.dispatch.region') {
return this.parseDispatchRegionOp(parser, result);
}
if (op === 'flow.dispatch.tensor.load' || op === 'flow.dispatch.tensor.store') {
return this.parseTensorLoadStoreOp(parser, result);
}
// Handle operations with visibility + symbol that aren't in schema or need manual parsing
if (op === 'flow.dispatch.entry') {
this.parseSymbolVisibility(parser, result);
const symName = parser.parseOptionalSymbolName();
if (symName !== null) {
result.addAttribute('sym_name', symName);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'flow.func') {
return this.parseFlowFuncOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseFlowFuncOp(parser, result) {
parser.parseOptionalVisibilityKeyword(result.attributes);
parser.parseSymbolName('sym_name', result.attributes);
const argResult = parser.parseFunctionArgumentList();
const inputs = argResult.arguments.map((a) => a.type);
const results = [];
if (parser.parseOptionalArrow()) {
const hasParens = parser.parseOptionalLParen();
if (!hasParens || !parser.parseOptionalRParen()) {
do {
const tiedOp = parser.parseOptionalOperand();
if (tiedOp) {
if (parser.parseOptionalKeyword('as')) {
const resultType = parser.parseType();
results.push(resultType);
} else {
results.push(new _.Type('tied'));
}
} else {
const resultType = parser.parseType();
results.push(resultType);
}
if (parser.parseOptionalLBrace()) {
const resultAttrs = new Map();
if (!parser.parseOptionalRBrace()) {
do {
const key = parser.parseAttribute();
resultAttrs.set(String(key.value || key), true);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
if (!hasParens) {
break;
}
} while (parser.parseOptionalComma());
if (hasParens) {
parser.parseRParen();
}
}
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion());
return true;
}
parseDispatchRegionOp(parser, result) {
const indexType = new _.IndexType();
const workloadOperands = parser.parseOperandList('optionalSquare');
const dynamicDimOperands = [];
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const type = parser.parseType();
if (parser.parseOptionalLBrace()) {
while (!parser.parseOptionalRBrace()) {
dynamicDimOperands.push(parser.parseOperand());
parser.parseOptionalComma();
}
}
result.addTypes([type]);
parser.parseOptionalComma();
}
}
}
parser.resolveOperands(dynamicDimOperands, dynamicDimOperands.map(() => indexType), result.operands);
parser.resolveOperands(workloadOperands, workloadOperands.map(() => indexType), result.operands);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
this.parseDispatchWorkgroupsCountRegion(parser, result);
return true;
}
parseDispatchWorkgroupsOp(parser, result) {
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
parser.parseKeyword(); // read subscript value
parser.parseOptionalComma();
}
}
const unresolvedOperands = parser.parseOperandList('paren');
if (parser.parseOptionalColon()) {
if (parser.parseOptionalLParen()) {
const inputTypes = parser.parseTypeList();
parser.parseRParen();
parser.resolveOperands(unresolvedOperands, inputTypes, result.operands);
const resultTypes = parser.parseArrowTypeList();
result.addTypes(resultTypes);
} else {
const types = parser.parseTypeList();
parser.resolveOperands(unresolvedOperands, types, result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
}
} else {
result.addTypes(parser.parseOptionalArrowTypeList());
}
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
if (parser.parseOptionalEqual()) {
const args = [];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const arg = parser.parseArgument(true, false);
args.push(arg);
parser.parseOptionalComma();
}
}
if (parser.parseOptionalArrow() || parser.parseOptionalKeyword('to')) {
parser.parseType();
}
const region = result.addRegion();
parser.parseRegion(region, args, /* enableNameShadowing */ true);
}
return true;
}
parseShapedFunctionType(parser, op, unresolvedArguments /*, otherArgs */) {
// unresolvedArguments is ref($arguments) - array of unresolved operands with .name and .number
const operandTypes = [];
if (parser.parseOptionalLParen()) {
let index = 0;
if (!parser.parseOptionalRParen()) {
do {
const type = parser.parseType();
if (type) {
operandTypes.push(type);
const startIdx = Math.max(0, op.operands.length - (index + 1));
if (startIdx + index < op.operands.length && !op.operands[startIdx + index].type) {
op.operands[startIdx + index].type = type;
}
index++;
}
if (parser.parseOptionalLBrace()) {
while (!parser.parseOptionalRBrace()) {
parser.parseOperand();
parser.parseOptionalComma();
}
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalArrow()) {
let index = 0;
const hasParens = parser.parseOptionalLParen();
if (!hasParens || !parser.parseOptionalRParen()) {
do {
const tiedResult = parser.parseOptionalOperand();
if (tiedResult) {
// Handle optional "as type" for tied results
if (parser.parseOptionalKeyword('as')) {
const type = parser.parseType();
if (type) {
if (index < op.types.length) {
op.types[index] = type;
} else {
op.addTypes([type]);
}
}
} else {
// Look up type from tied operand using unresolvedArguments
// unresolvedArguments contains UnresolvedOperand objects with .name and .number
const operands = unresolvedArguments || [];
const tiedOperandIndex = parser.findTiedOperand(tiedResult, operands);
if (tiedOperandIndex >= 0 && tiedOperandIndex < operandTypes.length) {
const type = operandTypes[tiedOperandIndex];
if (type) {
if (index < op.types.length) {
op.types[index] = type;
} else {
op.addTypes([type]);
}
}
}
}
index++;
} else {
const type = parser.parseType();
if (type) {
if (index < op.types.length) {
op.types[index] = type;
} else {
op.addTypes([type]);
}
index++;
}
}
if (parser.parseOptionalLBrace()) {
while (!parser.parseOptionalRBrace()) {
parser.parseOperand();
parser.parseOptionalComma();
}
}
if (!hasParens) {
break;
}
} while (parser.parseOptionalComma());
}
if (hasParens) {
parser.parseRParen();
}
}
}
parseTensorLoadStoreOp(parser, result) {
const op = result.name.getStringRef();
const isLoad = op.endsWith('.load');
const unresolvedOperands = [];
let nextOp = parser.parseOptionalOperand();
while (nextOp) {
unresolvedOperands.push(nextOp);
if (!parser.parseOptionalComma()) {
break;
}
nextOp = parser.parseOptionalOperand();
if (!nextOp) {
break;
}
}
const parseDynamicIndexList = (name) => {
parser.parseKeyword(name);
parser.parseEqual();
parser.parseLSquare();
const staticValues = [];
if (!parser.parseOptionalRSquare()) {
do {
const dynOp = parser.parseOptionalOperand();
if (dynOp) {
parser.resolveOperand(dynOp, null, result.operands);
staticValues.push(-9223372036854775808n);
} else {
const intVal = parser.parseOptionalInteger('int64');
if (intVal !== null) {
staticValues.push(intVal);
}
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute(`static_${name}`, staticValues);
};
parseDynamicIndexList('offsets');
parser.parseComma();
parseDynamicIndexList('sizes');
parser.parseComma();
parseDynamicIndexList('strides');
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const sourceType = parser.parseType();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => sourceType), result.operands);
if (parser.parseOptionalLBrace()) {
if (!parser.parseOptionalRBrace()) {
do {
const dim = parser.parseOperand();
parser.resolveOperand(dim, null, result.operands);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
if (parser.parseOptionalArrow() || parser.parseOptionalKeyword('to')) {
const resultType = parser.parseType();
if (isLoad && resultType) {
result.addTypes([resultType]);
}
}
return true;
}
parseDispatchWorkgroupBody(parser, op /*, args */) {
parser.parseLParen();
const regionArgs = [];
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
regionArgs.push({ name: arg, type: argType });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
const region = { blocks: [{ arguments: regionArgs, operations: [] }] };
parser.parseRegion(region);
op.regions.push(region);
}
parseDispatchWorkgroupsCountRegion(parser, op /*, args */) {
if (!parser.parseOptionalKeyword('count')) {
return;
}
parser.parseLParen();
const regionArgs = [];
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
regionArgs.push({ name: arg, type: argType });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseArrow();
if (parser.parseOptionalLParen()) {
parser.parseType();
parser.parseOptionalComma();
parser.parseType();
parser.parseOptionalComma();
parser.parseType();
parser.parseRParen();
} else {
parser.parseType();
parser.parseOptionalComma();
parser.parseType();
parser.parseOptionalComma();
parser.parseType();
}
const region = { blocks: [{ arguments: regionArgs, operations: [] }] };
parser.parseRegion(region);
op.regions.push(region);
}
parseShapedOperandList(parser, op, valuesOperands, valuesTypes, valueDimsOperands) {
do {
valuesOperands.push(parser.parseOperand());
parser.parseColon();
const valueType = parser.parseType();
valuesTypes.push(valueType);
if (valueType && valueType.shape) {
const dynamicDimCount = valueType.shape.filter((d) => d < 0n).length;
if (dynamicDimCount > 0 && parser.parseOptionalLBrace()) {
for (let i = 0; i < dynamicDimCount; i++) {
if (i > 0) {
parser.parseOptionalComma();
}
valueDimsOperands.push(parser.parseOperand());
}
parser.parseRBrace();
}
}
} while (parser.parseOptionalComma());
}
};
_.StreamDialect = class extends _.IREEDialect {
constructor(operations, name = 'stream') {
super(operations, name);
this.registerCustomDirective('DispatchOperands', this.parseDispatchOperands.bind(this));
this.registerCustomDirective('DispatchResources', this.parseDispatchResources.bind(this));
this.registerCustomDirective('ExplicitResourceRegion', this.parseExplicitResourceRegion.bind(this));
this.registerCustomDirective('ShapedTypeList', this.parseShapedTypeList.bind(this));
this.registerCustomDirective('ResourceRegion', this.parseResourceRegion.bind(this));
this.registerCustomDirective('ParameterLoadOperations', this.parseParameterLoadOperations.bind(this));
this.registerCustomDirective('EncodedResourceOperands', this.parseEncodedResourceOperands.bind(this));
this.registerCustomDirective('DispatchEntryPoints', this.parseDispatchEntryPoints.bind(this));
this.registerCustomDirective('ShapedTiedResult', this.parseShapedTiedResult.bind(this));
this.registerCustomDirective('EncodedShapedFunctionType', this.parseEncodedShapedFunctionType.bind(this));
this.registerCustomDirective('CollectiveParam', this.parseCollectiveParam.bind(this));
this.registerCustomDirective('PackSliceRanges', this.parsePackSliceRanges.bind(this));
this.registerCustomDirective('WorkgroupCountRegion', this.parseWorkgroupCountRegion.bind(this));
this.registerCustomDirective('DispatchFunctionSignature', this.parseDispatchFunctionSignature.bind(this));
this.registerCustomDirective('ShapedFunctionSignature', this.parseShapedFunctionSignature.bind(this));
this.registerCustomDirective('ConstantValueList', this.parseConstantValueList.bind(this));
this.registerCustomDirective('CmdCallOperands', this.parseCmdCallOperands.bind(this));
this.registerCustomDirective('ParameterReference', this.parseParameterReference.bind(this));
this.registerCustomDirective('ParameterGatherOperations', this.parseParameterGatherOperations.bind(this));
this.registerCustomDirective('ParameterScatterOperations', this.parseParameterScatterOperations.bind(this));
this.registerCustomDirective('SymbolAlias', this.parseSymbolAlias.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'resource') {
parser.parseLess();
const lifetime = parser.parseOptionalStar() ? '*' : parser.parseKeyword();
parser.parseGreater();
return new _.Type(`!${dialect}.resource<${lifetime}>`);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseDispatchResources(parser, op /*, args */) {
do {
const accessMode = parser.parseKeyword();
const unresolvedResource = parser.parseOperand();
parser.parseLSquare();
const unresolvedOffset = parser.parseOperand();
parser.parseKeyword('for');
const unresolvedLength = parser.parseOperand();
parser.parseRSquare();
parser.parseColon();
const resourceType = parser.parseType();
if (parser.parseOptionalLBrace()) {
while (!parser.parseOptionalRBrace()) {
parser.parseOperand();
parser.parseOptionalComma();
}
}
op.addAttribute('resource_access', accessMode);
parser.resolveOperand(unresolvedResource, resourceType, op.operands);
const indexType = new _.IndexType();
parser.resolveOperand(unresolvedOffset, indexType, op.operands);
parser.resolveOperand(unresolvedLength, indexType, op.operands);
} while (parser.parseOptionalComma());
}
parseShapedTypeList(parser, op, ...args) {
// 2-arg: custom<ShapedTypeList>(type($operands), $sizes)
// 3-arg: custom<ShapedTypeList>(type($operands), type($results), $sizes)
let operandTypes = null;
let resultTypes = null;
let sizeOperands = null;
if (args.length === 2) {
[operandTypes, sizeOperands] = args;
} else if (args.length >= 3) {
[operandTypes, resultTypes, sizeOperands] = args;
} else {
operandTypes = args[0] || null;
}
const indexType = new _.IndexType();
do {
const type = parser.parseType();
if (operandTypes) {
operandTypes.push(type);
}
if (resultTypes) {
resultTypes.push(type);
}
if (parser.parseOptionalLBrace()) {
do {
const sizeOperand = parser.parseOptionalOperand();
if (sizeOperand) {
if (sizeOperands) {
sizeOperands.push(sizeOperand);
}
parser.resolveOperand(sizeOperand, indexType, op.operands);
}
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
} while (parser.parseOptionalComma());
}
parseExplicitResourceRegion(parser, op /*, args */) {
parser.parseLParen();
const regionArgs = [];
const unresolvedOperands = [];
const operandTypes = [];
const unresolvedSizes = [];
if (!parser.parseOptionalRParen()) {
do {
const _optOp = parser.parseOptionalOperand();
if (_optOp) {
unresolvedOperands.push(_optOp);
}
parser.parseKeyword('as');
const arg = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
operandTypes.push(argType);
regionArgs.push({ name: arg, type: argType });
if (parser.parseOptionalLBrace()) {
const _sizeOp = parser.parseOptionalOperand();
if (_sizeOp) {
unresolvedSizes.push(_sizeOp);
}
parser.parseRBrace();
}
} while (parser.parseOptionalComma());
}
parser.parseRParen();
for (let i = 0; i < unresolvedOperands.length; i++) {
parser.resolveOperand(unresolvedOperands[i], operandTypes[i] || null, op.operands);
}
const indexType = new _.IndexType();
for (const unresolved of unresolvedSizes) {
parser.resolveOperand(unresolved, indexType, op.operands);
}
const region = { blocks: [{ arguments: regionArgs, operations: [] }] };
parser.parseRegion(region);
op.regions.push(region);
}
parseResourceRegion(parser, op /*, args */) {
const regionArgs = [];
const unresolvedOperands = [];
const operandTypes = [];
const unresolvedSizes = [];
const indexType = new _.IndexType();
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
parser.parseKeyword('as');
const arg = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
operandTypes.push(argType);
regionArgs.push({ name: arg, type: argType });
if (parser.parseOptionalLBrace()) {
const _sizeOp = parser.parseOptionalOperand();
if (_sizeOp) {
unresolvedSizes.push(_sizeOp);
}
parser.parseRBrace();
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
for (let i = 0; i < unresolvedOperands.length; i++) {
parser.resolveOperand(unresolvedOperands[i], operandTypes[i], op.operands);
}
for (const unresolved of unresolvedSizes) {
parser.resolveOperand(unresolved, indexType, op.operands);
}
const resultSizes = [];
const parseResultTypeOrTied = () => {
const tiedOp = parser.parseOptionalOperand();
if (tiedOp) {
if (parser.parseOptionalKeyword('as')) {
const resultType = parser.parseType();
op.addTypes([resultType]);
} else {
op.addTypes([new _.Type('tied')]);
}
} else {
const resultType = parser.parseType();
op.addTypes([resultType]);
}
if (parser.parseOptionalLBrace()) {
const sizeOp = parser.parseOptionalOperand();
if (sizeOp) {
resultSizes.push(sizeOp);
}
parser.parseRBrace();
}
};
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
parseResultTypeOrTied();
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
parseResultTypeOrTied();
}
}
for (const unresolved of resultSizes) {
parser.resolveOperand(unresolved, indexType, op.operands);
}
{
const region = { blocks: [{ arguments: regionArgs, operations: [] }] };
parser.parseRegion(region);
op.regions.push(region);
}
}
parseParameterLoadOperations(parser, op, sourceScopeOps, sourceKeysOps, sourceOffsetsOps, resultTypes, resultSizesOps) {
do {
const firstOperand = parser.parseOperand();
let key = firstOperand;
if (parser.parseOptionalColon()) {
sourceScopeOps.push(firstOperand);
parser.parseColon();
key = parser.parseOperand();
}
sourceKeysOps.push(key);
parser.parseLSquare();
sourceOffsetsOps.push(parser.parseOperand());
parser.parseRSquare();
parser.parseColon();
const resultType = parser.parseType();
op.addTypes([resultType]);
if (parser.parseOptionalLBrace()) {
resultSizesOps.push(parser.parseOperand());
parser.parseRBrace();
}
} while (parser.parseOptionalComma());
}
parseEncodedResourceOperands(parser, op) {
const indexType = new _.IndexType();
do {
parser.parseOperand();
parser.parseColon();
parser.parseType();
const encodingDims = parser.parseOperandList('optionalBrace');
parser.resolveOperands(encodingDims, encodingDims.map(() => indexType), op.operands);
parser.parseKeyword('in');
parser.parseType();
const sizeDims = parser.parseOperandList('optionalBrace');
parser.resolveOperands(sizeDims, sizeDims.map(() => indexType), op.operands);
} while (parser.parseOptionalComma());
}
parseDispatchEntryPoints(parser, op /*, args */) {
if (parser.parseOptionalLBrace()) {
do {
const symbol = parser.parseOptionalSymbolName();
op.addAttribute('entry_point', symbol);
} while (parser.parseOptionalComma());
parser.parseRBrace();
} else {
const symbol = parser.parseOptionalSymbolName();
op.addAttribute('entry_point', symbol);
}
}
parseShapedTiedResult(parser, op /*, args */) {
const tiedOp = parser.parseOptionalOperand();
if (tiedOp) {
parser.parseKeyword('as');
}
const type = parser.parseType();
op.types.push(type);
if (parser.parseOptionalLBrace()) {
const unresolvedSize = parser.parseOptionalOperand();
if (unresolvedSize) {
const indexType = new _.IndexType();
parser.resolveOperand(unresolvedSize, indexType, op.operands);
}
parser.parseRBrace();
}
}
parseEncodedShapedTypeList(parser, types, operands) {
const indexType = new _.IndexType();
do {
const type0 = parser.parseType();
const dims0 = parser.parseOperandList('optionalBrace');
if (operands) {
parser.resolveOperands(dims0, dims0.map(() => indexType), operands);
}
if (parser.parseOptionalKeyword('in')) {
const type1 = parser.parseType();
const dims1 = parser.parseOperandList('optionalBrace');
if (operands) {
parser.resolveOperands(dims1, dims1.map(() => indexType), operands);
}
types.push(type1);
} else {
types.push(type0);
}
} while (parser.parseOptionalComma());
}
parseEncodedShapedResultList(parser, operands, operandTypes, resultTypes, resolvedOperands) {
const indexType = new _.IndexType();
do {
let type0 = null;
const tryTied0 = parser.parseOptionalOperand();
if (!tryTied0) {
type0 = parser.parseType();
const dims0 = parser.parseOperandList('optionalBrace');
if (resolvedOperands) {
parser.resolveOperands(dims0, dims0.map(() => indexType), resolvedOperands);
}
}
if (!parser.parseOptionalKeyword('in')) {
if (type0) {
resultTypes.push(type0);
}
continue;
}
let resultType = null;
const tryTied1 = tryTied0 || parser.parseOptionalOperand();
if (tryTied1) {
const tiedOperandIndex = parser.findTiedOperand(tryTied1, operands);
if (parser.parseOptionalKeyword('as')) {
resultType = parser.parseType();
} else if (tiedOperandIndex >= 0 && operandTypes[tiedOperandIndex]) {
resultType = operandTypes[tiedOperandIndex];
}
} else {
resultType = parser.parseType();
}
const dims1 = parser.parseOperandList('optionalBrace');
if (resolvedOperands) {
parser.resolveOperands(dims1, dims1.map(() => indexType), resolvedOperands);
}
if (resultType) {
resultTypes.push(resultType);
}
} while (parser.parseOptionalComma());
}
parseEncodedShapedFunctionType(parser, op, operandsRef, operandTypes /* , ... */) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
this.parseEncodedShapedTypeList(parser, operandTypes, op.operands);
parser.parseRParen();
}
parser.parseArrow();
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
this.parseEncodedShapedResultList(parser, operandsRef, operandTypes, op.types, op.operands);
parser.parseRParen();
}
} else {
this.parseEncodedShapedResultList(parser, operandsRef, operandTypes, op.types, op.operands);
}
}
parseCollectiveParam(parser, op /*, args */) {
for (const keyword of ['source', 'target', 'source_target_pair']) {
if (parser.parseOptionalKeyword(keyword)) {
parser.parseLParen();
const unresolvedParam = parser.parseOperand();
parser.resolveOperand(unresolvedParam, null, op.operands);
parser.parseRParen();
return;
}
}
}
parsePackSliceRanges(parser, op, lifetimeIntervalsName, dynamicSliceSizesOperands, packedOffsetsTypes) {
const indexType = new _.IndexType();
while (parser.parseOptionalLSquare()) {
parser.parseAttribute();
parser.parseComma();
parser.parseAttribute();
parser.parseRSquare();
parser.parseEqual();
const unresolvedOperand = parser.parseOperand();
parser.resolveOperand(unresolvedOperand, indexType, op.operands);
packedOffsetsTypes.push(indexType);
if (!parser.parseOptionalComma()) {
break;
}
}
}
parseWorkgroupCountRegion(parser, op /*, workgroupCountName */) {
if (!parser.parseOptionalKeyword('workgroups')) {
return;
}
const region = { blocks: [] };
const block = { arguments: [], operations: [] };
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
if (parser.parseOptionalColon()) {
arg.type = parser.parseType();
}
block.arguments.push(arg);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalArrow()) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
parser.parseType();
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
region.blocks.push(block);
parser.parseRegion(region);
op.regions.push(region);
}
parseDispatchFunctionSignature(parser, op /*, args */) {
const indexType = new _.IndexType();
const inputs = [];
const results = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
parser.parseOperand();
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
parser.parseOperand();
if (parser.parseOptionalKeyword('for')) {
parser.parseOperand();
}
parser.parseRSquare();
}
}
parser.parseColon();
const type = parser.parseType();
inputs.push(type);
const dims = parser.parseOperandList('optionalBrace');
parser.resolveOperands(dims, dims.map(() => indexType), op.operands);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
const parseResultTypeOrTied = () => {
const tiedOp = parser.parseOptionalOperand();
if (tiedOp) {
if (parser.parseOptionalKeyword('as')) {
return parser.parseType();
}
return new _.Type('tied');
}
return parser.parseType();
};
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
results.push(parseResultTypeOrTied());
const dims = parser.parseOperandList('optionalBrace');
parser.resolveOperands(dims, dims.map(() => indexType), op.operands);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
results.push(parseResultTypeOrTied());
const dims = parser.parseOperandList('optionalBrace');
parser.resolveOperands(dims, dims.map(() => indexType), op.operands);
}
}
op.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
}
parseShapedFunctionSignature(parser, op /*, args */) {
this.parseDispatchFunctionSignature(parser, op);
}
parseConstantValueList(parser, op /*, args */) {
do {
const resultType = parser.parseType();
op.addTypes([resultType]);
if (parser.parseOptionalLBrace()) {
// Size is an SSA value like %c4, not an attribute
const _sizeOp = parser.parseOptionalOperand();
if (_sizeOp) {
parser.resolveOperand(_sizeOp, null, op.operands);
} else {
const size = parser.parseAttribute();
if (size) {
// If it's an integer literal, create operand directly
op.operands.push(new _.Value(null, size));
}
}
parser.parseRBrace();
}
parser.parseEqual();
parser.parseAttribute();
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
}
parseCmdCallOperands(parser, op /*, args */) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
const indexType = new _.IndexType();
do {
// Check for access mode keyword (ro, rw, wo)
const accessMode = parser.parseOptionalKeyword('ro') || parser.parseOptionalKeyword('rw') || parser.parseOptionalKeyword('wo');
if (accessMode) {
// Resource operand with offset/length: access operand[offset for length]
const unresolvedResource = parser.parseOperand();
parser.parseLSquare();
const unresolvedOffset = parser.parseOperand();
parser.parseKeyword('for');
const unresolvedLength = parser.parseOperand();
parser.parseRSquare();
op.addAttribute('resource_access', accessMode);
parser.resolveOperand(unresolvedResource, null, op.operands);
parser.resolveOperand(unresolvedOffset, indexType, op.operands);
parser.resolveOperand(unresolvedLength, indexType, op.operands);
} else {
// Primitive/custom operand
const unresolvedOperand = parser.parseOperand();
parser.resolveOperand(unresolvedOperand, null, op.operands);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parseParameterReference(parser, op, scopeOperands, keyOperands) {
const firstOperand = parser.parseOperand();
if (parser.parseOptionalColon()) {
parser.parseColon();
const keyOperand = parser.parseOperand();
if (scopeOperands) {
scopeOperands.push(firstOperand);
}
if (keyOperands) {
keyOperands.push(keyOperand);
}
} else if (keyOperands) {
keyOperands.push(firstOperand);
}
}
parseParameterGatherOperations(parser /*, op, args */) {
do {
// %scope::%key[offset] -> %target[offset for length] : type{size}
this.parseParameterReference(parser);
if (parser.parseOptionalLSquare()) {
parser.parseAttribute();
parser.parseRSquare();
}
parser.parseArrow();
parser.parseOperand();
if (parser.parseOptionalLSquare()) {
parser.parseAttribute();
parser.parseKeyword('for');
parser.parseAttribute();
parser.parseRSquare();
}
parser.parseColon();
parser.parseType();
if (parser.parseOptionalLBrace()) {
parser.parseOperand();
parser.parseRBrace();
}
} while (parser.parseOptionalComma());
}
parseParameterScatterOperations(parser /*, op, args */) {
do {
// %source[offset for length] : type{size} -> %scope::%key[offset]
parser.parseOperand();
if (parser.parseOptionalLSquare()) {
parser.parseAttribute();
parser.parseKeyword('for');
parser.parseAttribute();
parser.parseRSquare();
}
parser.parseColon();
parser.parseType();
if (parser.parseOptionalLBrace()) {
parser.parseOperand();
parser.parseRBrace();
}
parser.parseArrow();
this.parseParameterReference(parser);
if (parser.parseOptionalLSquare()) {
parser.parseAttribute();
parser.parseRSquare();
}
} while (parser.parseOptionalComma());
}
parseSymbolAlias(parser, op /*, args */) {
parser.parseSymbolName('sym_name', op.attributes);
if (parser.parseOptionalEqual()) {
const ref = parser.parseOptionalSymbolName();
op.addAttribute('function_ref', ref);
}
}
parseDispatchOperands(parser, op, resourceOperands /* offsets, ends, lengths */) {
// args are: [$resource_operands, $resource_operand_offsets, $resource_operand_ends, $resource_operand_lengths]
// resourceOperands is passed by ref so ShapedFunctionType can use it for tied operand lookup
parser.parseLParen();
if (parser.parseOptionalRParen()) {
return;
}
const unresolvedOperands = [];
do {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
if (Array.isArray(resourceOperands)) {
resourceOperands.push(operand);
}
if (parser.parseOptionalLSquare()) {
unresolvedOperands.push(parser.parseOperand()); // offset
parser.parseKeyword('to');
unresolvedOperands.push(parser.parseOperand()); // end
parser.parseKeyword('for');
unresolvedOperands.push(parser.parseOperand()); // length
parser.parseRSquare();
}
} while (parser.parseOptionalComma());
parser.parseRParen();
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, op.operands);
}
}
};
_.IOParametersDialect = class extends _.StreamDialect {
constructor(operations) {
super(operations, 'io_parameters');
}
};
_.PCFDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'pcf');
this.registerCustomDirective('ParallelExecutionBody', this.parseParallelExecutionBody.bind(this));
this.registerCustomDirective('InferNumIndexArgs', this.parseInferNumIndexArgs.bind(this));
}
parseInferNumIndexArgs() {
}
parseParallelExecutionBody(parser, op, initsOperands, initsTypes, dynamicSizesOperands, resultsTypes, isTiedAttr, regionName, numLeadingArgsAttr, parseOptionalLeadingArgs) {
const isTied = [];
const regionLeadingArgs = [];
const regionRefArgs = [];
const indexArgs = [];
if (parseOptionalLeadingArgs === '"true"') {
if (parser.parseOptionalArrow()) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
regionLeadingArgs.push({ value: arg, type: argType });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
op.addAttribute(numLeadingArgsAttr, regionLeadingArgs.length);
}
parser.parseKeyword('execute');
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const refArg = parser.parseOperand();
regionRefArgs.push({ value: refArg });
if (parser.parseOptionalEqual()) {
const initOperand = parser.parseOperand();
initsOperands.push(initOperand);
isTied.push(true);
} else {
isTied.push(false);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
const indexArg = parser.parseOperand();
parser.parseColon();
const indexType = parser.parseType();
indexArgs.push({ value: indexArg, type: indexType });
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (regionRefArgs.length > 0) {
parser.parseColon();
parser.parseLParen();
let refIdx = 0;
if (!parser.parseOptionalRParen()) {
do {
const refType = parser.parseType();
if (refIdx < regionRefArgs.length) {
regionRefArgs[refIdx].type = refType;
}
refIdx++;
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseArrow();
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
for (let i = 0; i < isTied.length; i++) {
if (i > 0) {
parser.parseComma();
}
const resType = parser.parseType();
resultsTypes.push(resType);
if (isTied[i]) {
initsTypes.push(resType);
} else if (parser.parseOptionalLBrace()) {
if (!parser.parseOptionalRBrace()) {
do {
dynamicSizesOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
}
parser.parseRParen();
}
}
if (isTied.length > 0) {
op.addAttribute(isTiedAttr, isTied);
}
const region = { blocks: [{ arguments: [...regionLeadingArgs, ...regionRefArgs, ...indexArgs], operations: [] }] };
parser.parseRegion(region);
op.regions.push(region);
}
};
_.IREEVectorExtDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'iree_vector_ext');
}
};
_.IREETensorExtDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'iree_tensor_ext');
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'dispatch.tensor') {
return _.iree.DispatchTensorType.parse(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.LinalgDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'linalg');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
const opInfo = result.name.getRegisteredInfo();
if (op === 'linalg.generic') {
return this.parseGenericOp(parser, result);
}
if (op === 'linalg.init_tensor') {
if (parser.parseOptionalLSquare()) {
const dims = [];
if (!parser.parseOptionalRSquare()) {
do {
const dimOp = parser.parseOptionalOperand();
if (dimOp) {
dims.push(dimOp.name);
} else {
const intVal = parser.parseOptionalInteger();
if (intVal !== null) {
dims.push(String(intVal));
}
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('static_sizes', dims);
}
result.addTypes(parser.parseOptionalColonTypeList());
return true;
}
if (op === 'linalg.fill') {
if (parser.parseOptionalLParen()) {
// Legacy format: linalg.fill(%operands) : types -> types
const unresolvedOperands = parser.parseOperandList();
parser.parseRParen();
parser.parseOptionalAttrDict(result.attributes);
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
return this.parseNamedStructuredOp(parser, result);
}
if (op === 'linalg.conv') {
let unresolvedOperands = [];
if (parser.parseOptionalLParen()) {
unresolvedOperands = parser.parseOperandList();
parser.parseRParen();
}
parser.parseOptionalAttrDict(result.attributes);
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
return true;
}
if (op === 'linalg.yield') {
const unresolvedOperands = parser.parseOperandList();
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
return true;
}
if (op === 'linalg.pack') {
return this.parsePackOp(parser, result);
}
if (op === 'linalg.unpack') {
return this.parseUnpackOp(parser, result);
}
if (op === 'linalg.transpose') {
return this.parseDstStyleOp(parser, result, (parser, attributes) => {
parser.parseDenseI64ArrayAttr('permutation', attributes);
});
}
if (op === 'linalg.reduce') {
// Optional short form: { payload_op attr-dict }
let payloadOpName = null;
const payloadOpAttrs = new Map();
if (parser.parseOptionalLBrace()) {
payloadOpName = parser.parseCustomOperationName();
parser.parseOptionalAttrDict(payloadOpAttrs);
parser.parseRBrace();
}
if (!this.parseDstStyleOp(parser, result, (parser, attributes) => {
parser.parseDenseI64ArrayAttr('dimensions', attributes);
})) {
return false;
}
// Parse block arguments and region (or add body with payload op)
if (payloadOpName) {
this.addBodyWithPayloadOp(result, payloadOpName, payloadOpAttrs, true, true);
} else {
const regionArgs = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const value = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
regionArgs.push({ value, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
}
return true;
}
if (op === 'linalg.broadcast') {
return this.parseDstStyleOp(parser, result, (parser, attributes) => {
parser.parseDenseI64ArrayAttr('dimensions', attributes);
});
}
if (op === 'linalg.elementwise') {
parser.parseKeyword('kind');
parser.parseEqual();
const kind = parser.parseAttribute();
result.addAttribute('kind', kind.value);
const indexingMapsAttr = this.parseIndexingMapsAttr(parser);
if (indexingMapsAttr !== null) {
result.addAttribute('indexing_maps', indexingMapsAttr);
}
return this.parseNamedStructuredOp(parser, result);
}
if (op === 'linalg.map') {
return this.parseMapOp(parser, result);
}
if (op === 'linalg.contract') {
const indexingMapsAttr = this.parseIndexingMapsAttr(parser);
if (!indexingMapsAttr) {
parser.emitError("Expected 'indexing_maps' attribute");
}
result.addAttribute('indexing_maps', indexingMapsAttr);
return this.parseNamedStructuredOp(parser, result);
}
if (op === 'linalg.abs' || op === 'linalg.add' ||
op === 'linalg.batch_matmul' || op === 'linalg.batch_matvec' || op === 'linalg.batch_mmt4d' || op === 'linalg.batch_reduce_matmul' || op === 'linalg.batch_vecmat' ||
op === 'linalg.ceil' || op === 'linalg.copy' ||
op === 'linalg.conv_1d' || op === 'linalg.conv_1d_ncw_fcw' || op === 'linalg.conv_1d_nwc_wcf' ||
op === 'linalg.conv_2d' || op === 'linalg.conv_2d_nchw_fchw' || op === 'linalg.conv_2d_nchw_fchw_q' || op === 'linalg.conv_2d_ngchw_fgchw' || op === 'linalg.conv_2d_ngchw_gfchw' || op === 'linalg.conv_2d_ngchw_gfchw_q' || op === 'linalg.conv_2d_nhwc_fhwc' || op === 'linalg.conv_2d_nhwc_fhwc_q' || op === 'linalg.conv_2d_nhwc_hwcf' || op === 'linalg.conv_2d_nhwc_hwcf_q' || op === 'linalg.conv_2d_nhwgc_gfhwc' || op === 'linalg.conv_2d_nhwgc_gfhwc_q' ||
op === 'linalg.conv_3d' || op === 'linalg.conv_3d_ncdhw_fcdhw' || op === 'linalg.conv_3d_ndhwc_dhwcf' || op === 'linalg.conv_3d_ndhwc_dhwcf_q' ||
op === 'linalg.depthwise_conv_1d_ncw_cw' || op === 'linalg.depthwise_conv_1d_nwc_wc' || op === 'linalg.depthwise_conv_1d_nwc_wcm' ||
op === 'linalg.depthwise_conv_2d_nchw_chw' || op === 'linalg.depthwise_conv_2d_nhwc_hwc' || op === 'linalg.depthwise_conv_2d_nhwc_hwc_q' || op === 'linalg.depthwise_conv_2d_nhwc_hwcm' || op === 'linalg.depthwise_conv_2d_nhwc_hwcm_q' ||
op === 'linalg.depthwise_conv_3d_ncdhw_cdhw' || op === 'linalg.depthwise_conv_3d_ndhwc_dhwc' || op === 'linalg.depthwise_conv_3d_ndhwc_dhwcm' ||
op === 'linalg.div' || op === 'linalg.div_unsigned' || op === 'linalg.dot' ||
op === 'linalg.erf' || op === 'linalg.exp' ||
op === 'linalg.fill_rng_2d' || op === 'linalg.floor' ||
op === 'linalg.log' ||
op === 'linalg.matmul' || op === 'linalg.matmul_transpose_a' || op === 'linalg.matmul_transpose_b' || op === 'linalg.matvec' || op === 'linalg.max' || op === 'linalg.min' || op === 'linalg.mmt4d' || op === 'linalg.mul' ||
op === 'linalg.negf' ||
op === 'linalg.pooling_nchw_max' || op === 'linalg.pooling_nchw_sum' || op === 'linalg.pooling_ncw_max' || op === 'linalg.pooling_ncw_sum' || op === 'linalg.pooling_ndhwc_max' || op === 'linalg.pooling_ndhwc_min' || op === 'linalg.pooling_ndhwc_sum' || op === 'linalg.pooling_nhwc_max' || op === 'linalg.pooling_nhwc_max_unsigned' || op === 'linalg.pooling_nhwc_min' || op === 'linalg.pooling_nhwc_min_unsigned' || op === 'linalg.pooling_nhwc_sum' || op === 'linalg.pooling_nwc_max' || op === 'linalg.pooling_nwc_max_unsigned' || op === 'linalg.pooling_nwc_min' || op === 'linalg.pooling_nwc_min_unsigned' || op === 'linalg.pooling_nwc_sum' || op === 'linalg.powf' ||
op === 'linalg.quantized_batch_matmul' || op === 'linalg.quantized_matmul' ||
op === 'linalg.reciprocal' || op === 'linalg.round' || op === 'linalg.rsqrt' ||
op === 'linalg.select' || op === 'linalg.sqrt' || op === 'linalg.square' || op === 'linalg.sub' ||
op === 'linalg.tanh' ||
op === 'linalg.vecmat') {
const indexingMapsAttr = this.parseIndexingMapsAttr(parser);
if (indexingMapsAttr) {
result.addAttribute('indexing_maps', indexingMapsAttr);
}
return this.parseNamedStructuredOp(parser, result);
}
if (opInfo.metadata && opInfo.metadata.assemblyFormat) {
return super.parseOperation(parser, result);
}
const parsed = this.parseCommonStructuredOpParts(parser, result);
if (!parsed) {
return false;
}
if (parser.parseOptionalKeyword('attrs')) {
parser.parseEqual();
parser.parseOptionalAttrDict(result.attributes);
}
result.addTypes(parser.parseOptionalArrowTypeList());
parser.parseOptionalRegion(result.addRegion(), []);
return true;
}
// "Common parsing used for both named structured ops created by ods-gen and by
// manually defined C++ ops. Does not handle regions."
// Returns { inputTypes, outputTypes } on success, null on failure.
parseCommonStructuredOpParts(parser, result, addOperandSegmentSizes = true) {
if (parser.parseOptionalLess()) {
result.propertiesAttr = parser.parseAttribute();
parser.parseGreater();
}
parser.parseOptionalAttrDict(result.attributes);
const inputTypes = [];
const outputTypes = [];
if (parser.parseOptionalKeyword('ins')) {
if (!parser.parseOptionalLParen()) {
return null;
}
const unresolvedIns = [];
let insOp = parser.parseOptionalOperand();
while (insOp) {
unresolvedIns.push(insOp);
if (!parser.parseOptionalComma()) {
break;
}
insOp = parser.parseOptionalOperand();
}
if (parser.parseOptionalColon()) {
if (!parser.parseOptionalRParen()) {
do {
inputTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.resolveOperands(unresolvedIns, inputTypes, result.operands);
parser.parseRParen();
}
} else if (!parser.parseOptionalRParen()) {
return null;
}
}
if (parser.parseOptionalKeyword('outs')) {
if (!parser.parseOptionalLParen()) {
return null;
}
const unresolvedOuts = [];
let outsOp = parser.parseOptionalOperand();
while (outsOp) {
unresolvedOuts.push(outsOp);
if (!parser.parseOptionalComma()) {
break;
}
outsOp = parser.parseOptionalOperand();
}
if (parser.parseOptionalColon()) {
if (!parser.parseOptionalRParen()) {
do {
outputTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.resolveOperands(unresolvedOuts, outputTypes, result.operands);
parser.parseRParen();
}
} else if (!parser.parseOptionalRParen()) {
return null;
}
}
if (addOperandSegmentSizes) {
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([inputTypes.length, outputTypes.length]));
}
return { inputTypes, outputTypes };
}
parseNamedStructuredOp(parser, result) {
const parsed = this.parseCommonStructuredOpParts(parser, result);
if (!parsed) {
return false;
}
parser.parseOptionalAttrDict(result.attributes);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
parseIndexingMapsAttr(parser) {
if (!parser.parseOptionalKeyword('indexing_maps')) {
return null;
}
parser.parseEqual();
return parser.parseAttribute();
}
parseDstStyleOp(parser, op, parseAttrsFn) {
const parsed = this.parseCommonStructuredOpParts(parser, op, false);
if (!parsed) {
return false;
}
for (const outputType of parsed.outputTypes) {
if (outputType instanceof _.RankedTensorType) {
op.addTypes([outputType]);
}
}
if (parseAttrsFn) {
parseAttrsFn(parser, op.attributes);
}
parser.parseOptionalAttrDict(op.attributes);
return true;
}
addBodyWithPayloadOp(op, payloadOpName, payloadOpAttrs /*, initFirst, mapInit */) {
const region = op.addRegion();
const block = { operations: [], arguments: [] };
for (const operand of op.operands) {
if (operand && operand.type) {
const elemType = operand.type.elementType || operand.type;
block.arguments.push({ value: null, type: elemType });
}
}
const payloadState = new _.OperationState(null,payloadOpName);
if (op.operands.length > 0) {
const lastOperand = op.operands[op.operands.length - 1];
if (lastOperand && lastOperand.type) {
const elemType = lastOperand.type.elementType || lastOperand.type;
payloadState.addTypes([elemType]);
}
}
for (const [name, value] of payloadOpAttrs) {
payloadState.attributes.set(name, value);
}
block.operations.push(_.Operation.create(payloadState));
const yieldState = new _.OperationState(null, this.getOperation('linalg.yield'));
block.operations.push(_.Operation.create(yieldState));
region.blocks = [block];
}
parseMapOp(parser, result) {
let payloadOpName = null;
const payloadOpAttrs = new Map();
if (parser.parseOptionalLBrace()) {
payloadOpName = parser.parseCustomOperationName();
parser.parseOptionalAttrDict(payloadOpAttrs);
parser.parseRBrace();
}
if (!this.parseDstStyleOp(parser, result)) {
return false;
}
// Parse block arguments and region (or add body with payload op)
if (payloadOpName) {
if (result.operands.length > 0) {
this.addBodyWithPayloadOp(result, payloadOpName, payloadOpAttrs, false, false);
} else {
result.addRegion();
}
} else {
const regionArgs = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const value = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
regionArgs.push({ value, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
}
return true;
}
parseGenericOp(parser, result) {
const optAttr = parser.parseOptionalAttribute();
if (optAttr === null) {
parser.parseOptionalAttrDict(result.attributes);
} else {
result.addAttribute('trait', optAttr);
}
this.parseCommonStructuredOpParts(parser, result);
if (parser.parseOptionalKeyword('attrs')) {
parser.parseEqual();
parser.parseOptionalAttrDict(result.attributes);
}
const region = result.addRegion();
parser.parseRegion(region);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
parsePackOp(parser, result) {
const unresolvedSource = parser.parseOperand();
let unresolvedPadding = null;
let paddingType = null;
if (parser.parseOptionalKeyword('padding_value')) {
parser.parseLParen();
unresolvedPadding = parser.parseOperand();
parser.parseColon();
paddingType = parser.parseType();
parser.parseRParen();
}
if (parser.parseOptionalKeyword('outer_dims_perm')) {
parser.parseEqual();
const outerDimsPerm = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
outerDimsPerm.push(parser.parseInteger('int64'));
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('outer_dims_perm', outerDimsPerm);
}
parser.parseKeyword('inner_dims_pos');
parser.parseEqual();
const innerDimsPos = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
innerDimsPos.push(parser.parseInteger('int64'));
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('inner_dims_pos', innerDimsPos);
parser.parseKeyword('inner_tiles');
parser.parseEqual();
const staticInnerTiles = [];
const dynamicTileOperands = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
const dynOp = parser.parseOptionalOperand();
if (dynOp) {
dynamicTileOperands.push(dynOp);
staticInnerTiles.push(_.ShapedType.kDynamic);
} else {
staticInnerTiles.push(parser.parseInteger('int64'));
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('static_inner_tiles', staticInnerTiles);
parser.parseKeyword('into');
const unresolvedDest = parser.parseOperand();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const sourceType = parser.parseType();
parser.parseArrow();
const destType = parser.parseType();
parser.resolveOperand(unresolvedSource, sourceType, result.operands);
parser.resolveOperand(unresolvedDest, destType, result.operands);
if (unresolvedPadding) {
parser.resolveOperand(unresolvedPadding, paddingType, result.operands);
}
for (const dynOp of dynamicTileOperands) {
parser.resolveOperand(dynOp, null, result.operands);
}
if (!(sourceType instanceof _.MemRefType)) {
result.addTypes([destType]);
}
return true;
}
parseUnpackOp(parser, result) {
const unresolvedSource = parser.parseOperand();
if (parser.parseOptionalKeyword('outer_dims_perm')) {
parser.parseEqual();
const outerDimsPerm = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
outerDimsPerm.push(parser.parseInteger('int64'));
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('outer_dims_perm', outerDimsPerm);
}
parser.parseKeyword('inner_dims_pos');
parser.parseEqual();
const innerDimsPos = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
innerDimsPos.push(parser.parseInteger('int64'));
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('inner_dims_pos', innerDimsPos);
parser.parseKeyword('inner_tiles');
parser.parseEqual();
const staticInnerTiles = [];
const dynamicTileOperands = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
const dynOp = parser.parseOptionalOperand();
if (dynOp) {
dynamicTileOperands.push(dynOp);
staticInnerTiles.push(-9223372036854775808n); // ShapedType::kDynamic
} else {
staticInnerTiles.push(parser.parseInteger('int64'));
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('static_inner_tiles', staticInnerTiles);
parser.parseKeyword('into');
const unresolvedDest = parser.parseOperand();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const sourceType = parser.parseType();
parser.parseArrow();
const destType = parser.parseType();
parser.resolveOperand(unresolvedSource, sourceType, result.operands);
parser.resolveOperand(unresolvedDest, destType, result.operands);
for (const dynOp of dynamicTileOperands) {
parser.resolveOperand(dynOp, null, result.operands);
}
if (!(sourceType instanceof _.MemRefType)) {
result.addTypes([destType]);
}
return true;
}
};
_.ONNXDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'onnx');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'onnx.Constant') {
parser.parseOptionalAttrDict(result.attributes);
if (result.attributes.size === 0) {
const value = parser.parseOptionalAttribute();
if (value) {
const name = value instanceof _.SparseElementsAttr ? 'sparse_value' : 'value';
result.addAttribute(name, value);
result.addTypes([value.type]);
return true;
}
}
result.addTypes([parser.parseColonType()]);
return true;
}
if (op === 'onnx.ConstantOfShape') {
parser.parseLParen();
const unresolved = parser.parseOperand();
parser.parseRParen();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
parser.parseLParen();
const inputType = parser.parseType();
parser.resolveOperand(unresolved, inputType, result.operands);
parser.parseRParen();
parser.parseArrow();
const outputType = parser.parseType();
result.addTypes([outputType]);
return true;
}
return super.parseOperation(parser, result);
}
};
_.KrnlDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'krnl');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'krnl.define_loops') {
const count = parser.parseOptionalInteger();
if (count !== null) {
result.addAttribute('num_loops', parseInt(count, 10));
const loopType = new _.Type('!krnl.loop');
const types = Array(count).fill(loopType);
result.addTypes(types);
}
return true;
}
if (op === 'krnl.get_linear_offset_index') {
const unresolvedOperands = [];
const staticIndices = [];
const memref = parser.parseOperand();
unresolvedOperands.push(memref);
if (parser.parseOptionalKeyword('at')) {
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
// Indices can be either SSA values (%arg) or integer constants (0, 10, etc.)
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedOperands.push(operand);
staticIndices.push(-9223372036854775808n); // ShapedType::kDynamic marker
} else {
staticIndices.push(parser.parseInteger('int64'));
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (staticIndices.length > 0) {
result.addAttribute('static_indices', staticIndices);
}
}
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, type, result.operands);
}
result.addTypes([new _.IndexType()]);
return true;
}
if (op === 'krnl.prefetch') {
const unresolvedOperands = [];
const memref = parser.parseOperand();
unresolvedOperands.push(memref);
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const index = parser.parseOperand();
unresolvedOperands.push(index);
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
parser.parseComma();
const readOrWrite = parser.parseKeyword();
result.addAttribute('isWrite', readOrWrite === 'write');
parser.parseComma();
parser.parseKeyword('locality');
parser.parseLess();
const localityHint = parser.parseInteger();
result.addAttribute('localityHint', localityHint);
parser.parseGreater();
parser.parseComma();
const cacheType = parser.parseKeyword();
result.addAttribute('isDataCache', cacheType === 'data');
parser.parseOptionalAttrDict(result.attributes);
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, type, result.operands);
}
return true;
}
if (op === 'krnl.iterate') {
const unresolvedOperands = parser.parseOperandList('paren');
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
result.addAttribute('num_optimized_loops', result.operands.length);
const indexType = new _.IndexType();
const boundMaps = [];
const inductionVarRefs = [];
parser.parseKeyword('with');
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
parser.parseCommaSeparatedList('none', () => {
const loop = parser.parseOperand();
parser.resolveOperand(loop, null, result.operands);
parser.parseArrow();
const inductionVar = parser.parseOperand();
inductionVarRefs.push(inductionVar);
parser.parseEqual();
for (let i = 0; i < 2; i++) {
if (i > 0) {
parser.parseKeyword('to');
}
parser.parseOptionalKeyword(i === 0 ? 'max' : 'min');
const boundAttr = parser.parseAttribute();
if (boundAttr instanceof _.AffineMapAttr || (boundAttr && boundAttr.value instanceof _.AffineMap)) {
boundMaps.push(boundAttr);
parser.parseDimAndSymbolList(result.operands);
} else {
const op = parser.parseOptionalOperand();
if (op) {
parser.resolveOperand(op, indexType, result.operands);
}
boundMaps.push(boundAttr);
}
}
});
parser.parseRParen();
}
result.addAttribute('bounds', boundMaps);
const entryArgs = inductionVarRefs.map((ref) => ({ value: ref, type: indexType }));
if (parser.parseOptionalKeyword('iter_args')) {
const iterOperands = [];
const iterArgs = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
iterArgs.push(parser.parseOperand());
parser.parseEqual();
iterOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
const iterTypes = parser.parseOptionalArrowTypeList();
result.addTypes(iterTypes);
for (let i = 0; i < iterArgs.length; i++) {
const type = i < iterTypes.length ? iterTypes[i] : null;
if (type) {
parser.resolveOperand(iterOperands[i], type, result.operands);
}
entryArgs.push({ value: iterArgs[i], type });
}
}
const region = {};
parser.parseRegion(region, entryArgs);
result.regions = [region];
return true;
}
return super.parseOperation(parser, result);
}
};
_.hlo.MhloDialect = class extends _.hlo.HLODialect {
constructor(operations) {
super(operations, 'mhlo');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
const opInfo = result.name.getRegisteredInfo();
if (opInfo.metadata && opInfo.metadata.parser && opInfo.metadata.parser.includes('parseOneResultSameOperandTypeOp')) {
return this.parseOneResultSameOperandTypeOp(parser, result);
}
if (op === 'mhlo.constant') {
if (parser.parseOptionalLParen() && parser.parseOptionalRParen()) {
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
parser.parseLParen();
parser.parseRParen();
parser.parseArrow();
const type = parser.parseType();
result.addTypes([type]);
}
} else {
parser.parseOptionalAttrDict(result.attributes);
const value = parser.parseAttribute();
if (value) {
result.addAttribute('value', value);
if (value.type && result.types.length === 0) {
result.addTypes([value.type]);
}
}
if (result.types.length === 0) {
result.addTypes(parser.parseOptionalColonTypeList());
}
}
return true;
}
if (op === 'mhlo.reduce') {
return super.parseReduceOp(parser, result, (dims) => dims, 'mhlo.return');
}
if (op === 'mhlo.scan') {
return super.parseScanOp(parser, result, 'mhlo.return');
}
if (op === 'mhlo.while') {
return super.parseWhileOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseOneResultSameOperandTypeOp(parser, result) {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
const types = unresolvedOperands.map(() => type);
parser.resolveOperands(unresolvedOperands, types, result.operands);
if (result.types.length > 0) {
result.types[0] = type;
}
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
return true;
}
};
_.hlo.ChloDialect = class extends _.hlo.HLODialect {
constructor(operations) {
super(operations, 'chlo');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'chlo.scan') {
return super.parseScanOp(parser, result, 'stablehlo.return');
}
return super.parseOperation(parser, result);
}
};
_.THLODialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'thlo');
}
parseOperation(parser, result) {
const opInfo = result.name.getRegisteredInfo();
if (opInfo.metadata.assemblyFormat) {
return super.parseOperation(parser, result);
}
if (parser.parseOptionalKeyword('ins')) {
parser.parseLParen();
for (let insOp = parser.parseOptionalOperand(); insOp; insOp = parser.parseOptionalOperand()) {
const operand = insOp;
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRParen();
}
if (parser.parseOptionalKeyword('outs')) {
parser.parseLParen();
for (let outsOp = parser.parseOptionalOperand(); outsOp; outsOp = parser.parseOptionalOperand()) {
const operand = outsOp;
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRParen();
}
parser.parseOptionalAttrDict(result.attributes);
const blockArguments = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const value = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
blockArguments.push({ value, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
if (parser.parseOptionalLBrace()) {
const region = { blocks: [] };
const block = { operations: [], arguments: blockArguments };
while (!parser.parseOptionalRBrace()) {
const operation = parser.parseOperation();
block.operations.push(_.Operation.create(operation));
}
region.blocks.push(block);
result.regions.push(region);
}
return true;
}
};
_.quant = {};
_.quant.UniformQuantizedType = class extends _.Type {
constructor(storageType, storageRange, expressedType, scale, zeroPoint) {
super(null);
this.storageType = storageType;
this.storageRange = storageRange;
this.expressedType = expressedType;
this.scale = scale;
this.zeroPoint = zeroPoint;
}
toString() {
const range = this.storageRange ? `<${this.storageRange[0]}:${this.storageRange[1]}>` : '';
const zp = this.zeroPoint === 0 ? '' : `:${this.zeroPoint}`;
return `!quant.uniform<${this.storageType}${range}:${this.expressedType}, ${this.scale}${zp}>`;
}
};
_.quant.UniformQuantizedPerAxisType = class extends _.Type {
constructor(storageType, storageRange, expressedType, quantizedDimension, scales, zeroPoints) {
super(null);
this.storageType = storageType;
this.storageRange = storageRange;
this.expressedType = expressedType;
this.quantizedDimension = quantizedDimension;
this.scales = scales;
this.zeroPoints = zeroPoints;
}
toString() {
const range = this.storageRange ? `<${this.storageRange[0]}:${this.storageRange[1]}>` : '';
const params = this.scales.map((s, i) => {
const zp = this.zeroPoints[i] === 0 ? '' : `:${this.zeroPoints[i]}`;
return `${s}${zp}`;
}).join(', ');
return `!quant.uniform<${this.storageType}${range}:${this.expressedType}:${this.quantizedDimension}, {${params}}>`;
}
};
_.quant.AnyQuantizedType = class extends _.Type {
constructor(storageType, storageRange, expressedType) {
super(null);
this.storageType = storageType;
this.storageRange = storageRange;
this.expressedType = expressedType;
}
toString() {
const range = this.storageRange ? `<${this.storageRange[0]}:${this.storageRange[1]}>` : '';
const expressed = this.expressedType ? `:${this.expressedType}` : '';
return `!quant.any<${this.storageType}${range}${expressed}>`;
}
};
_.quant.CalibratedQuantizedType = class extends _.Type {
constructor(expressedType, min, max) {
super(null);
this.expressedType = expressedType;
this.min = min;
this.max = max;
}
toString() {
return `!quant.calibrated<${this.expressedType}<${this.min}:${this.max}>>`;
}
};
_.quant.QuantileQuantizedType = class extends _.Type {
constructor(storageType, quantileType, quantiles, storageRange) {
super(null);
this.storageType = storageType;
this.quantileType = quantileType;
this.quantiles = quantiles;
this.storageRange = storageRange;
}
toString() {
const range = this.storageRange ? `, <${this.storageRange[0]}:${this.storageRange[1]}>` : '';
return `!quant.quantile<${this.storageType}:${this.quantileType}, {${this.quantiles.join(',')}}${range}>`;
}
};
_.quant.QuantDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'quant');
}
parseType(parser) {
const typeNameSpelling = parser.parseKeyword();
if (typeNameSpelling === 'uniform') {
return this.parseUniformType(parser);
}
if (typeNameSpelling === 'any') {
return this.parseAnyType(parser);
}
if (typeNameSpelling === 'calibrated') {
return this.parseCalibratedType(parser);
}
if (typeNameSpelling === 'quantile') {
return this.parseQuantileType(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown quantized type '${typeNameSpelling}'`);
return null;
}
parseQuantileType(parser) {
parser.parseLess();
const storageType = parser.parseType();
parser.parseColon();
const quantileType = parser.parseType();
parser.parseComma();
parser.parseLBrace();
const quantiles = [];
if (!parser.parseOptionalRBrace()) {
do {
quantiles.push(parser.parseAttribute(new _.FloatType('f64')).value);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
let storageRange = null;
if (parser.parseOptionalComma()) {
parser.parseLess();
const min = parser.parseInteger();
parser.parseColon();
const max = parser.parseInteger();
parser.parseGreater();
storageRange = [min, max];
}
parser.parseGreater();
return new _.quant.QuantileQuantizedType(storageType, quantileType, quantiles, storageRange);
}
parseUniformType(parser) {
parser.parseLess();
let storageType = parser.parseOptionalType();
if (!storageType) {
const keyword = parser.parseKeyword();
storageType = new _.Type(keyword);
}
let storageRange = null;
if (parser.parseOptionalLess()) {
const min = parser.parseInteger();
parser.parseColon();
const max = parser.parseInteger();
parser.parseGreater();
storageRange = [min, max];
}
parser.parseColon();
const expressedType = parser.parseType();
let quantizedDimension = null;
let isSubChannel = false;
if (parser.parseOptionalColon()) {
const dim = parser.parseOptionalInteger();
if (dim === null) {
isSubChannel = true;
parser.parseLBrace();
if (!parser.parseOptionalRBrace()) {
do {
parser.parseInteger();
parser.parseColon();
parser.parseInteger();
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
} else {
quantizedDimension = dim;
}
}
parser.parseComma();
const parseQuantParams = () => {
const attr = parser.parseAttribute(new _.FloatType('f64'));
let zeroPoint = 0;
if (parser.parseOptionalColon()) {
zeroPoint = parser.parseInteger();
}
return { scale: attr.value, zeroPoint };
};
const isPerTensor = quantizedDimension === null && !isSubChannel;
if (isPerTensor) {
const { scale, zeroPoint } = parseQuantParams();
parser.parseGreater();
return new _.quant.UniformQuantizedType(storageType, storageRange, expressedType, scale, zeroPoint);
}
const scales = [];
const zeroPoints = [];
const parseQuantParamList = () => {
do {
if (parser.parseOptionalLBrace()) {
parseQuantParamList();
} else {
const { scale, zeroPoint } = parseQuantParams();
scales.push(scale);
zeroPoints.push(zeroPoint);
}
} while (parser.parseOptionalComma());
parser.parseRBrace();
};
parser.parseLBrace();
parseQuantParamList();
parser.parseGreater();
if (isSubChannel) {
const range = storageRange ? `<${storageRange[0]}:${storageRange[1]}>` : '';
const params = scales.map((s, i) => {
const zp = zeroPoints[i] === 0 ? '' : `:${zeroPoints[i]}`;
return `${s}${zp}`;
}).join(', ');
return new _.Type(`!quant.uniform<${storageType}${range}:${expressedType}:{}, {${params}}>`);
}
return new _.quant.UniformQuantizedPerAxisType(storageType, storageRange, expressedType, quantizedDimension, scales, zeroPoints);
}
parseAnyType(parser) {
parser.parseLess();
let storageType = parser.parseOptionalType();
if (!storageType) {
const keyword = parser.parseKeyword();
storageType = new _.Type(keyword);
}
let storageRange = null;
if (parser.parseOptionalLess()) {
const min = parser.parseInteger();
parser.parseColon();
const max = parser.parseInteger();
parser.parseGreater();
storageRange = [min, max];
}
let expressedType = null;
if (parser.parseOptionalColon()) {
expressedType = parser.parseType();
}
parser.parseGreater();
return new _.quant.AnyQuantizedType(storageType, storageRange, expressedType);
}
parseCalibratedType(parser) {
parser.parseLess();
const expressedType = parser.parseType();
parser.parseLess();
const min = parser.parseAttribute(new _.FloatType('f64')).value;
parser.parseColon();
const max = parser.parseAttribute(new _.FloatType('f64')).value;
parser.parseGreater();
parser.parseGreater();
return new _.quant.CalibratedQuantizedType(expressedType, min, max);
}
};
_.TosaDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tosa');
this._enumOps = new Map([
['tosa.rescale', 'rounding_mode'],
['tosa.apply_scale', 'rounding_mode'],
['tosa.resize', 'mode'],
['tosa.argmax', 'nan_mode'],
['tosa.max_pool2d', 'nan_mode'],
['tosa.max_pool2d_adaptive', 'nan_mode'],
['tosa.clamp', 'nan_mode'],
['tosa.maximum', 'nan_mode'],
['tosa.minimum', 'nan_mode'],
['tosa.reduce_max', 'nan_mode'],
['tosa.reduce_min', 'nan_mode'],
['tosa.matmul_t_block_scaled', 'block_size'],
['tosa.cast_from_block_scaled', 'block_size'],
['tosa.cast_to_block_scaled', 'block_size'],
['tosa.conv2d_block_scaled', 'block_size'],
]);
this.registerCustomDirective('VariableOpTypeOrInitialValue', this.parseVariableOpTypeOrInitialValue.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'shape') {
parser.parseLess();
const rank = parser.parseInteger();
parser.parseGreater();
return new _.Type(`!${dialect}.shape<${rank}>`);
}
if (mnemonic === 'mxint8') {
return new _.Type(`!${dialect}.mxint8`);
}
if (mnemonic === 'block_scaled') {
parser.parseLess();
const blockShape = parser.parseOptionalKeyword();
parser.parseColon();
const scaleType = parser.parseType();
parser.parseColon();
const valueType = parser.parseType();
let scaleValues = '';
if (parser.parseOptionalComma()) {
const values = parser.parseCommaSeparatedList('brace', () => parser.parseAttribute());
scaleValues = `, {${values.join(', ')}}`;
}
parser.parseGreater();
return new _.Type(`!${dialect}.block_scaled<${blockShape}:${scaleType}:${valueType}${scaleValues}>`);
}
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tosa.cond_if') {
return this.parseIfOp(parser, result);
}
if (op === 'tosa.while_loop') {
return this.parseWhileOp(parser, result);
}
if (this._enumOps.has(op)) {
return this.parseWithEnumHandling(parser, result, this._enumOps.get(op));
}
return super.parseOperation(parser, result);
}
parseIfOp(parser, result) {
const thenRegion = result.addRegion();
const elseRegion = result.addRegion();
const cond = parser.parseOperand();
const regionArgs = [];
const operands = [];
const hasAssignment = parser.parseOptionalAssignmentList(regionArgs, operands);
parser.parseColon();
const condType = parser.parseType();
parser.resolveOperand(cond, condType, result.operands);
if (hasAssignment) {
const functionType = parser.parseType();
result.addTypes(functionType.results);
parser.resolveOperands(operands, functionType.inputs, result.operands);
} else {
result.addTypes(parser.parseOptionalArrowTypeList());
}
parser.parseRegion(thenRegion);
if (parser.parseOptionalKeyword('else')) {
parser.parseRegion(elseRegion);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseWhileOp(parser, result) {
const regionArgs = [];
const operands = [];
parser.parseOptionalAssignmentList(regionArgs, operands);
const functionType = parser.parseColonType();
result.addTypes(functionType.results);
parser.resolveOperands(operands, functionType.inputs, result.operands);
for (let i = 0; i < regionArgs.length; i++) {
regionArgs[i].type = functionType.inputs[i];
}
const condRegion = result.addRegion();
parser.parseRegion(condRegion, regionArgs);
parser.parseKeyword('do');
const bodyRegion = result.addRegion();
parser.parseRegion(bodyRegion);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
parseAttrEntryWithEnumHandling(parser, attrs, enumKey) {
const name = parser.parseKeyword();
parser.parseEqual();
if (name === enumKey) {
const kw = parser.parseOptionalKeyword();
if (kw) {
attrs.set(name, kw);
return;
}
}
const attr = parser.parseAttribute();
attrs.set(name, attr);
}
parseWithEnumHandling(parser, result, enumKey) {
const operands = [];
do {
operands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
const attrs = new Map();
if (parser.parseOptionalLBrace()) {
if (!parser.parseOptionalRBrace()) {
do {
this.parseAttrEntryWithEnumHandling(parser, attrs, enumKey);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
parser.parseColon();
const fnTy = parser.parseType();
parser.resolveOperands(operands, fnTy.inputs, result.operands);
result.addTypes(fnTy.results);
for (const [name, value] of attrs) {
result.attributes.set(name, value);
}
return true;
}
parseVariableOpTypeOrInitialValue(parser, op /*, args */) {
if (parser.parseOptionalEqual()) {
const initialValue = parser.parseAttribute();
op.addAttribute('initial_value', initialValue);
} else if (parser.parseOptionalColon()) {
const type = parser.parseType();
op.addAttribute('type', type);
}
}
};
_.IRDLDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'irdl');
this.registerCustomDirective('SingleBlockRegion', this.parseSingleBlockRegion.bind(this));
this.registerCustomDirective('NamedValueList', this.parseNamedValueList.bind(this));
this.registerCustomDirective('NamedValueListWithVariadicity', this.parseNamedValueListWithVariadicity.bind(this));
this.registerCustomDirective('AttributesOp', this.parseAttributesOp.bind(this));
}
parseSingleBlockRegion(parser, op, /* args */) {
parser.parseOptionalRegion(op.addRegion());
}
parseNamedValueList(parser, op, argsAttrName, namesAttrName) {
const argValues = [];
const nameValues = [];
const parseOne = () => {
const name = parser.parseKeyword();
nameValues.push(name);
parser.parseColon();
const value = parser.parseOperand();
argValues.push(value);
return value;
};
parser.parseCommaSeparatedList('paren', parseOne);
if (argsAttrName && namesAttrName) {
// args contains SSA values (%0, %1, ...) - resolve and add as operands
for (const value of argValues) {
parser.resolveOperand(value, null, op.operands);
}
op.addAttribute(namesAttrName, nameValues);
}
}
parseNamedValueListWithVariadicity(parser, op, argsAttrName, namesAttrName, variadicityAttrName) {
const argValues = [];
const nameValues = [];
const variadicityValues = [];
const parseOne = () => {
const variadicity = parser.parseOptionalKeyword(['single', 'optional', 'variadic']);
const name = parser.parseKeyword();
nameValues.push(name);
parser.parseColon();
const value = parser.parseOperand();
argValues.push(value);
variadicityValues.push(variadicity || 'single');
return value;
};
parser.parseCommaSeparatedList('paren', parseOne);
if (argsAttrName && namesAttrName) {
// args contains SSA values (%0, %1, ...) - resolve and add as operands
for (const value of argValues) {
parser.resolveOperand(value, null, op.operands);
}
op.addAttribute(namesAttrName, nameValues);
if (variadicityAttrName) {
op.addAttribute(variadicityAttrName, variadicityValues);
}
}
}
parseAttributesOp(parser, op, argsAttrName, namesAttrName) {
const argValues = [];
const nameValues = [];
if (parser.parseOptionalLBrace()) {
if (!parser.parseOptionalRBrace()) {
do {
const name = parser.parseString();
nameValues.push(name);
parser.parseEqual();
const value = parser.parseOperand();
argValues.push(value);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
if (argsAttrName && namesAttrName) {
for (const value of argValues) {
parser.resolveOperand(value, null, op.operands);
}
op.addAttribute(namesAttrName, nameValues);
}
}
};
_.xegpu = {};
_.xegpu.TensorDescType = class extends _.Type {
constructor(shape, elementType, encoding, layout) {
super(null);
this.shape = shape;
this.elementType = elementType;
this.encoding = encoding || null;
this.layout = layout || null;
}
static parse(parser) {
parser.parseLess();
const shape = parser.parseDimensionList().dimensions;
const elementType = parser.parseType();
let encoding = null;
let layout = null;
while (parser.parseOptionalComma()) {
const attr = parser.parseAttribute();
if (attr) {
const attrStr = attr.toString ? attr.toString() : String(attr);
if (attrStr.includes('layout') || attrStr.includes('distribute')) {
layout = attr;
} else {
encoding = attr;
}
}
}
parser.parseGreater();
return new _.xegpu.TensorDescType(shape, elementType, encoding, layout);
}
toString() {
const dims = this.shape.map((d) => d === -1 ? '?' : String(d)).join('x');
let result = `!xegpu.tensor_desc<${dims}x${this.elementType}`;
if (this.encoding) {
result += `, ${this.encoding}`;
}
if (this.layout) {
result += `, ${this.layout}`;
}
result += '>';
return result;
}
};
_.xegpu.MemDescType = class extends _.Type {
constructor(shape, elementType, layout) {
super(null);
this.shape = shape;
this.elementType = elementType;
this.layout = layout || null;
}
static parse(parser) {
parser.parseLess();
const shape = parser.parseDimensionList(false, true).dimensions;
const elementType = parser.parseType();
let layout = null;
if (parser.parseOptionalComma()) {
layout = parser.parseAttribute();
}
parser.parseGreater();
return new _.xegpu.MemDescType(shape, elementType, layout);
}
toString() {
const dims = this.shape.map((d) => String(d)).join('x');
let result = `!xegpu.mem_desc<${dims}x${this.elementType}`;
if (this.layout) {
result += `, ${this.layout}`;
}
result += '>';
return result;
}
};
_.xegpu.XeGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xegpu');
this.registerCustomDirective('OptionalDynamicIndexList', this.parseOptionalDynamicIndexList.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'tensor_desc') {
return _.xegpu.TensorDescType.parse(parser);
}
if (mnemonic === 'mem_desc') {
return _.xegpu.MemDescType.parse(parser);
}
if (mnemonic === 'nbarrier') {
return new _.Type(`!${dialect}.nbarrier`);
}
}
return null;
}
parseOptionalDynamicIndexList(parser, op, dynamicAttrName, staticAttrName) {
const indices = [];
const dynamicValues = [];
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const operand = parser.parseOptionalOperand();
if (operand) {
dynamicValues.push(operand);
indices.push(-9223372036854775808n);
} else {
const intVal = parser.parseOptionalInteger('int64');
if (intVal === null) {
break;
}
indices.push(intVal);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (dynamicAttrName && staticAttrName) {
for (const value of dynamicValues) {
parser.resolveOperand(value, null, op.operands);
}
if (indices.length > 0) {
op.addAttribute(staticAttrName, indices);
}
}
}
}
};
_.ShardDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'shard');
this.registerCustomDirective('DimensionList', this.parseDimensionList.bind(this));
}
parseDimensionList(parser, op, attrName) {
if (parser.parseOptionalLSquare()) {
parser.parseOptionalRSquare();
op.addAttribute(attrName, []);
return;
}
const dimInfo = parser.parseDimensionList(true, false);
if (attrName) {
op.addAttribute(attrName, dimInfo.dimensions);
}
}
};
_.spirv = {};
_.spirv.RuntimeArrayType = class extends _.Type {
constructor(elementType, stride, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.elementType = elementType;
this.stride = stride || 0;
}
static parse(parser, dialect) {
parser.parseLess();
const elementType = parser.parseType();
let stride = 0;
if (parser.parseOptionalComma()) {
parser.parseKeyword('stride');
parser.parseEqual();
stride = parser.parseInteger();
}
parser.parseGreater();
return new _.spirv.RuntimeArrayType(elementType, stride, dialect);
}
toString() {
if (this.stride) {
return `!${this.dialect}.rtarray<${this.elementType}, stride=${this.stride}>`;
}
return `!${this.dialect}.rtarray<${this.elementType}>`;
}
};
_.spirv.StructType = class extends _.Type {
constructor(memberTypes, memberDecorations, identifier, structDecorations, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.memberTypes = memberTypes;
this.memberDecorations = memberDecorations;
this.identifier = identifier || null;
this.structDecorations = structDecorations || [];
}
static parse(parser, dialect) {
parser.parseLess();
let identifier = null;
const name = parser.parseOptionalKeyword();
if (name) {
identifier = name;
if (parser.parseOptionalGreater()) {
return new _.spirv.StructType([], [], identifier, undefined, dialect);
}
parser.parseComma();
}
parser.parseLParen();
const memberTypes = [];
const memberDecorations = [];
if (!parser.parseOptionalRParen()) {
for (;;) {
memberTypes.push(parser.parseType());
if (parser.parseOptionalLSquare()) {
const dec = [];
for (;;) {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
const kw = parser.parseOptionalKeyword();
if (kw) {
if (parser.parseOptionalEqual()) {
const val = parser.parseInteger();
dec.push(`${kw}=${val}`);
} else {
dec.push(kw);
}
}
} else {
dec.push(intVal);
}
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRSquare();
memberDecorations.push(dec);
} else {
memberDecorations.push(null);
}
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRParen();
}
const structDecorations = [];
while (parser.parseOptionalComma()) {
structDecorations.push(parser.parseKeyword());
}
parser.parseGreater();
const result = new _.spirv.StructType(memberTypes, memberDecorations, identifier, structDecorations, dialect);
return result;
}
toString() {
const members = this.memberTypes.map((t, i) => {
const dec = this.memberDecorations?.[i];
return dec ? `${t} [${dec.join(', ')}]` : `${t}`;
}).join(', ');
const decs = this.structDecorations.length > 0 ? `, ${this.structDecorations.join(', ')}` : '';
if (this.identifier && this.memberTypes.length === 0) {
return `!${this.dialect}.struct<${this.identifier}>`;
}
if (this.identifier) {
return `!${this.dialect}.struct<${this.identifier}, (${members})${decs}>`;
}
return `!${this.dialect}.struct<(${members})${decs}>`;
}
};
_.spirv.PointerType = class extends _.Type {
constructor(pointeeType, storageClass, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.pointeeType = pointeeType;
this.storageClass = storageClass;
}
static parse(parser, dialect) {
parser.parseLess();
const pointeeType = parser.parseType();
parser.parseComma();
const storageClass = parser.parseKeyword();
parser.parseGreater();
return new _.spirv.PointerType(pointeeType, storageClass, dialect);
}
toString() {
return `!${this.dialect}.ptr<${this.pointeeType}, ${this.storageClass}>`;
}
};
_.spirv.ArrayType = class extends _.Type {
constructor(elementType, count, stride, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.elementType = elementType;
this.count = count;
this.stride = stride;
}
static parse(parser, dialect) {
parser.parseLess();
const dims = parser.parseDimensionList(false, true);
const count = dims.dimensions[0];
const elementType = parser.parseType();
let stride = 0;
if (parser.parseOptionalComma()) {
parser.parseKeyword('stride');
parser.parseEqual();
stride = parser.parseInteger();
}
parser.parseGreater();
return new _.spirv.ArrayType(elementType, count, stride, dialect);
}
toString() {
const strideStr = this.stride ? `, stride=${this.stride}` : '';
return `!${this.dialect}.array<${this.count} x ${this.elementType}${strideStr}>`;
}
};
_.spirv.MatrixType = class extends _.Type {
constructor(columnType, columnCount, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.columnType = columnType;
this.columnCount = columnCount;
}
static parse(parser, dialect) {
parser.parseLess();
const dims = parser.parseDimensionList(false, true);
const columnCount = dims.dimensions[0];
const columnType = parser.parseType();
parser.parseGreater();
return new _.spirv.MatrixType(columnType, columnCount, dialect);
}
toString() {
return `!${this.dialect}.matrix<${this.columnCount} x ${this.columnType}>`;
}
};
_.spirv.SampledImageType = class extends _.Type {
constructor(imageType, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.imageType = imageType;
}
static parse(parser, dialect) {
parser.parseLess();
const imageType = parser.parseType();
parser.parseGreater();
return new _.spirv.SampledImageType(imageType, dialect);
}
toString() {
return `!${this.dialect}.sampled_image<${this.imageType}>`;
}
};
_.spirv.ImageType = class extends _.Type {
constructor(elementType, dim, depthInfo, arrayedInfo, samplingInfo, samplerUseInfo, imageFormat, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.elementType = elementType;
this.dim = dim;
this.depthInfo = depthInfo;
this.arrayedInfo = arrayedInfo;
this.samplingInfo = samplingInfo;
this.samplerUseInfo = samplerUseInfo;
this.imageFormat = imageFormat;
}
static parse(parser, dialect) {
parser.parseLess();
const elementType = parser.parseType();
parser.parseComma();
const dim = parser.parseKeyword();
parser.parseComma();
const depthInfo = parser.parseKeyword();
parser.parseComma();
const arrayedInfo = parser.parseKeyword();
parser.parseComma();
const samplingInfo = parser.parseKeyword();
parser.parseComma();
const samplerUseInfo = parser.parseKeyword();
parser.parseComma();
const imageFormat = parser.parseKeyword();
parser.parseGreater();
return new _.spirv.ImageType(elementType, dim, depthInfo, arrayedInfo, samplingInfo, samplerUseInfo, imageFormat, dialect);
}
toString() {
return `!${this.dialect}.image<${this.elementType}, ${this.dim}, ${this.depthInfo}, ${this.arrayedInfo}, ${this.samplingInfo}, ${this.samplerUseInfo}, ${this.imageFormat}>`;
}
};
_.spirv.CooperativeMatrixType = class extends _.Type {
constructor(rows, columns, elementType, scope, use, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.rows = rows;
this.columns = columns;
this.elementType = elementType;
this.scope = scope;
this.use = use;
}
static parse(parser, dialect) {
parser.parseLess();
const dims = parser.parseDimensionList(false, true);
const rows = dims.dimensions[0];
const columns = dims.dimensions[1];
const elementType = parser.parseType();
parser.parseComma();
const scope = parser.parseKeyword();
parser.parseComma();
const use = parser.parseKeyword();
parser.parseGreater();
return new _.spirv.CooperativeMatrixType(rows, columns, elementType, scope, use, dialect);
}
toString() {
return `!${this.dialect}.coopmatrix<${this.rows}x${this.columns}x${this.elementType}, ${this.scope}, ${this.use}>`;
}
};
_.spirv.ArmTensorType = class extends _.Type {
constructor(dims, elementType, dialect = 'spirv') {
super(null);
this.dialect = dialect;
this.dims = dims;
this.elementType = elementType;
}
static parse(parser, dialect) {
parser.parseLess();
let dims = null;
if (parser.parseOptionalStar()) {
dims = null;
parser.parseXInDimensionList();
} else {
dims = parser.parseDimensionList(true, true).dimensions;
}
const elementType = parser.parseType();
parser.parseGreater();
return new _.spirv.ArmTensorType(dims, elementType, dialect);
}
toString() {
if (this.dims === null) {
return `!${this.dialect}.arm.tensor<*x${this.elementType}>`;
}
const dimsStr = this.dims.map((d) => d < 0 ? '?' : d).join('x');
return `!${this.dialect}.arm.tensor<${dimsStr}x${this.elementType}>`;
}
};
_.spirv.SPIRVDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'spirv');
this.registerCustomType('SPIRV_Type', (parser) => parser.parseType());
this.registerCustomDirective('ImageOperands', this.parseImageOperands.bind(this));
this.registerCustomDirective('SwitchOpCases', this.parseSwitchOpCases.bind(this));
this.registerCustomDirective('SPIRV_I32_1DArmTensor', this.parseSPIRV_I32_1DArmTensor.bind(this));
this.registerCustomAttribute('SPIRV_ScopeAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_MemorySemanticsAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_MemoryAccessAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_GroupOperationAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_KHR_CooperativeMatrixLayoutAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_KHR_CooperativeMatrixOperandsAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('SPIRV_TosaExtNaNPropagationModeAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'ptr') {
return _.spirv.PointerType.parse(parser, dialect);
}
if (mnemonic === 'rtarray') {
return _.spirv.RuntimeArrayType.parse(parser, dialect);
}
if (mnemonic === 'struct') {
return _.spirv.StructType.parse(parser, dialect);
}
if (mnemonic === 'array') {
return _.spirv.ArrayType.parse(parser, dialect);
}
if (mnemonic === 'matrix') {
return _.spirv.MatrixType.parse(parser, dialect);
}
if (mnemonic === 'sampled_image') {
return _.spirv.SampledImageType.parse(parser, dialect);
}
if (mnemonic === 'image') {
return _.spirv.ImageType.parse(parser, dialect);
}
if (mnemonic === 'coopmatrix') {
return _.spirv.CooperativeMatrixType.parse(parser, dialect);
}
if (mnemonic === 'arm.tensor') {
return _.spirv.ArmTensorType.parse(parser, dialect);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parseSwitchOpCases(parser, op, selectorTypes, defaultTargetName, defaultOperands, defaultOperandTypes, literalsAttr, targetsName, targetOperands, targetOperandTypes) {
if (!parser.parseOptionalKeyword('default')) {
return;
}
if (!parser.parseOptionalColon()) {
return;
}
const defaultDestination = parser.parseOptionalSuccessor();
if (!defaultDestination) {
return;
}
const defaultDest = { label: defaultDestination, arguments: [] };
if (parser.parseOptionalLParen()) {
const unresolvedDefaults = parser.parseOperandList();
for (const value of unresolvedDefaults) {
defaultDest.arguments.push({ value });
defaultOperands.push(value);
}
const defTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < defTypes.length; idx++) {
defaultOperandTypes.push(defTypes[idx]);
if (idx < defaultDest.arguments.length) {
defaultDest.arguments[idx].type = defTypes[idx];
}
}
parser.parseOptionalRParen();
}
op.successors = op.successors || [];
op.successors.push(defaultDest);
const caseValues = [];
while (parser.parseOptionalComma()) {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
caseValues.push(value);
if (!parser.parseOptionalColon()) {
break;
}
const caseDestination = parser.parseOptionalSuccessor();
if (!caseDestination) {
break;
}
const caseDest = { label: caseDestination, arguments: [] };
if (parser.parseOptionalLParen()) {
const caseOps = parser.parseOperandList();
for (const argValue of caseOps) {
caseDest.arguments.push({ value: argValue });
targetOperands.push(argValue);
}
const caseTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < caseTypes.length; idx++) {
targetOperandTypes.push(caseTypes[idx]);
if (idx < caseDest.arguments.length) {
caseDest.arguments[idx].type = caseTypes[idx];
}
}
parser.parseOptionalRParen();
}
op.successors.push(caseDest);
}
if (caseValues.length > 0) {
op.addAttribute(literalsAttr, caseValues);
}
}
parseImageOperands(parser, op) {
if (parser.parseOptionalLSquare()) {
const value = parser.parseAttribute();
op.addAttribute('image_operands', value);
parser.parseRSquare();
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// spv.Constant value : type
if (op === 'spirv.Constant' || op === 'spv.Constant') {
const value = parser.parseAttribute();
result.addAttribute('value', value);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
} else if (value && value.type) {
result.addTypes([value.type]);
}
return true;
}
// SPIR-V comparison operations: result is i1 (or vector<Nxi1> for vector operands)
const spirvCompareOps = new Set([
'spirv.SLessThan', 'spv.SLessThan',
'spirv.SLessThanEqual', 'spv.SLessThanEqual',
'spirv.SGreaterThan', 'spv.SGreaterThan',
'spirv.SGreaterThanEqual', 'spv.SGreaterThanEqual',
'spirv.ULessThan', 'spv.ULessThan',
'spirv.ULessThanEqual', 'spv.ULessThanEqual',
'spirv.UGreaterThan', 'spv.UGreaterThan',
'spirv.UGreaterThanEqual', 'spv.UGreaterThanEqual',
'spirv.IEqual', 'spv.IEqual',
'spirv.INotEqual', 'spv.INotEqual',
'spirv.FOrdEqual', 'spv.FOrdEqual',
'spirv.FOrdNotEqual', 'spv.FOrdNotEqual',
'spirv.FOrdLessThan', 'spv.FOrdLessThan',
'spirv.FOrdLessThanEqual', 'spv.FOrdLessThanEqual',
'spirv.FOrdGreaterThan', 'spv.FOrdGreaterThan',
'spirv.FOrdGreaterThanEqual', 'spv.FOrdGreaterThanEqual',
'spirv.FUnordEqual', 'spv.FUnordEqual',
'spirv.FUnordNotEqual', 'spv.FUnordNotEqual',
'spirv.FUnordLessThan', 'spv.FUnordLessThan',
'spirv.FUnordLessThanEqual', 'spv.FUnordLessThanEqual',
'spirv.FUnordGreaterThan', 'spv.FUnordGreaterThan',
'spirv.FUnordGreaterThanEqual', 'spv.FUnordGreaterThanEqual'
]);
if (spirvCompareOps.has(op)) {
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
const operandType = parser.parseType();
parser.resolveOperands(unresolvedOperands, [operandType, operandType], result.operands);
// Result type is i1 for scalar, vector<Nxi1> for vector
if (operandType instanceof _.VectorType) {
result.addTypes([new _.VectorType(operandType.shape, new _.IntegerType('i1'), operandType.scalableDims)]);
} else {
result.addTypes([new _.IntegerType('i1')]);
}
} else {
parser.resolveOperands(unresolvedOperands, [null, null], result.operands);
result.addTypes([new _.IntegerType('i1')]);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'spirv.GL.FClamp' || op === 'spirv.GL.UClamp' || op === 'spirv.GL.SClamp' || op === 'spirv.GL.NClamp' || op === 'spirv.GL.Fma' || op === 'spirv.GL.SmoothStep' ||
op === 'spirv.GLSL.FClamp' || op === 'spirv.GLSL.UClamp' || op === 'spirv.GLSL.SClamp' || op === 'spirv.GLSL.NClamp' || op === 'spirv.GLSL.Fma' || op === 'spirv.GLSL.SmoothStep' ||
op === 'spv.GL.FClamp' || op === 'spv.GL.UClamp' || op === 'spv.GL.SClamp' || op === 'spv.GL.NClamp' || op === 'spv.GL.Fma' || op === 'spv.GL.SmoothStep' ||
op === 'spv.GLSL.FClamp' || op === 'spv.GLSL.UClamp' || op === 'spv.GLSL.SClamp' || op === 'spv.GLSL.NClamp' || op === 'spv.GLSL.Fma' || op === 'spv.GLSL.SmoothStep') {
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
const type = parser.parseType();
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, type, result.operands);
}
result.addTypes([type]);
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
return true;
}
if (op === 'spirv.SpecConstantComposite' || op === 'spv.SpecConstantComposite') {
parser.parseSymbolName('sym_name', result.attributes);
parser.parseLParen();
const constituents = [];
if (!parser.parseOptionalRParen()) {
do {
const sym = parser.parseOptionalSymbolName();
if (sym) {
constituents.push(sym);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
result.addAttribute('constituents', constituents);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addAttribute('type', type);
}
return true;
}
if (op === 'spirv.EXT.SpecConstantCompositeReplicate') {
parser.parseSymbolName('sym_name', result.attributes);
parser.parseLParen();
const constituent = parser.parseOptionalSymbolName();
if (constituent) {
result.addAttribute('constituent', constituent);
}
parser.parseRParen();
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addAttribute('type', type);
}
return true;
}
if (op === 'spirv.SpecConstantOperation' || op === 'spv.SpecConstantOperation') {
parser.parseKeyword('wraps');
const wrappedOp = parser.parseGenericOperation();
if (wrappedOp) {
const region = { blocks: [{ operations: [wrappedOp] }] };
result.regions.push(region);
if (wrappedOp.results && wrappedOp.results.length > 0) {
result.addTypes([wrappedOp.results[0].type]);
}
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'spirv.Load' || op === 'spv.Load') {
const storageClass = parser.parseString();
result.addAttribute('storage_class', storageClass);
const ptrOperand = parser.parseOperand();
if (parser.parseOptionalLSquare()) {
const memoryAccess = [];
if (!parser.parseOptionalRSquare()) {
do {
const str = parser.parseOptionalString();
if (str === null) {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
memoryAccess.push(String(intVal));
} else {
memoryAccess.push(str);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (memoryAccess.length > 0) {
result.addAttribute('memory_access', memoryAccess.join(', '));
}
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
const ptrType = new _.spirv.PointerType(type, storageClass);
parser.resolveOperand(ptrOperand, ptrType, result.operands);
result.addTypes([type]);
} else {
parser.resolveOperand(ptrOperand, null, result.operands);
}
return true;
}
if (op === 'spirv.CompositeExtract' || op === 'spv.CompositeExtract') {
const compositeOperand = parser.parseOperand();
if (parser.parseOptionalLSquare()) {
const indices = [];
if (!parser.parseOptionalRSquare()) {
do {
const intVal = parser.parseOptionalInteger();
if (intVal !== null) {
indices.push(intVal);
}
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('indices', indices);
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
parser.resolveOperand(compositeOperand, type, result.operands);
result.addTypes([type]);
} else {
parser.resolveOperand(compositeOperand, null, result.operands);
}
return true;
}
// Handle AccessChain with old syntax (no -> for result type)
// Uses compatibility flag because assemblyFormat includes `->` which is not present in old files
if (op === 'spirv.AccessChain' || op === 'spv.AccessChain') {
result.compatibility = true;
const unresolvedOperands = [];
unresolvedOperands.push(parser.parseOperand());
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
unresolvedOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const basePtrType = parser.parseType();
const types = [basePtrType];
while (parser.parseOptionalComma()) {
types.push(parser.parseType());
}
parser.resolveOperands(unresolvedOperands, types, result.operands);
if (parser.parseOptionalArrow()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
} else {
// Old syntax without explicit result type
result.addTypes([basePtrType]);
}
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
return true;
}
if (op === 'spirv.Variable' || op === 'spv.Variable') {
let unresolvedInit = null;
if (parser.parseOptionalKeyword('init')) {
parser.parseLParen();
unresolvedInit = parser.parseOperand();
parser.parseRParen();
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
if (unresolvedInit) {
// The init value type should match the element type of the pointer result
parser.resolveOperand(unresolvedInit, null, result.operands);
}
} else if (unresolvedInit) {
parser.resolveOperand(unresolvedInit, null, result.operands);
}
return true;
}
if (op === 'spirv.Store' || op === 'spv.Store') {
const storageClass = parser.parseString();
result.addAttribute('storage_class', storageClass);
const unresolvedOperands = [];
unresolvedOperands.push(parser.parseOperand());
parser.parseComma();
unresolvedOperands.push(parser.parseOperand());
if (parser.parseOptionalLSquare()) {
const memoryAccess = [];
if (!parser.parseOptionalRSquare()) {
do {
const str = parser.parseOptionalString();
if (str === null) {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
memoryAccess.push(String(intVal));
} else {
memoryAccess.push(str);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (memoryAccess.length > 0) {
result.addAttribute('memory_access', memoryAccess.join(', '));
}
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
const ptrType = new _.spirv.PointerType(type, storageClass);
parser.resolveOperands(unresolvedOperands, [ptrType, type], result.operands);
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
return true;
}
if (op === 'spirv.CompositeInsert' || op === 'spv.CompositeInsert') {
const unresolvedOperands = [];
unresolvedOperands.push(parser.parseOperand());
parser.parseComma();
unresolvedOperands.push(parser.parseOperand());
if (parser.parseOptionalLSquare()) {
const indices = [];
if (!parser.parseOptionalRSquare()) {
do {
const intVal = parser.parseOptionalInteger();
if (intVal !== null) {
indices.push(intVal);
}
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('indices', indices);
}
if (parser.parseOptionalColon()) {
const objType = parser.parseType();
parser.resolveOperand(unresolvedOperands[0], objType, result.operands);
if (parser.parseOptionalKeyword('into')) {
const compositeType = parser.parseType();
parser.resolveOperand(unresolvedOperands[1], compositeType, result.operands);
result.addTypes([compositeType]);
} else {
parser.resolveOperand(unresolvedOperands[1], null, result.operands);
}
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
return true;
}
if (op === 'spirv.BranchConditional' || op === 'spv.BranchConditional') {
const conditionOperand = parser.parseOperand();
parser.resolveOperand(conditionOperand, null, result.operands);
// Parse optional branch weights [trueWeight, falseWeight]
if (parser.parseOptionalLSquare()) {
const weights = [];
if (!parser.parseOptionalRSquare()) {
do {
const intVal = parser.parseOptionalInteger();
if (intVal !== null) {
weights.push(String(intVal));
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (weights.length > 0) {
result.addAttribute('branch_weights', weights);
}
}
parser.parseComma();
if (!result.successors) {
result.successors = [];
}
const trueLabel = parser.parseOptionalSuccessor();
const trueSucc = { label: trueLabel };
if (parser.parseOptionalLParen()) {
const trueOperands = parser.parseOperandList();
trueSucc.arguments = trueOperands.map((value) => ({ value }));
const trueTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < trueTypes.length && idx < trueSucc.arguments.length; idx++) {
trueSucc.arguments[idx].type = trueTypes[idx];
}
parser.parseRParen();
}
result.successors.push(trueSucc);
parser.parseComma();
const falseLabel = parser.parseOptionalSuccessor();
const falseSucc = { label: falseLabel };
if (parser.parseOptionalLParen()) {
const falseOperands = parser.parseOperandList();
falseSucc.arguments = falseOperands.map((value) => ({ value }));
const falseTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < falseTypes.length && idx < falseSucc.arguments.length; idx++) {
falseSucc.arguments[idx].type = falseTypes[idx];
}
parser.parseRParen();
}
result.successors.push(falseSucc);
return true;
}
if (op === 'spirv.CompositeConstruct' || op === 'spv.CompositeConstruct') {
result.compatibility = true;
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
if (parser.parseOptionalLParen()) {
const types = parser.parseTypeList();
parser.parseRParen();
parser.parseArrow();
parser.resolveOperands(unresolvedOperands, types, result.operands);
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
const type = parser.parseType();
result.addTypes([type]);
} else {
for (const unresolvedOp of unresolvedOperands) {
parser.resolveOperand(unresolvedOp, null, result.operands);
}
}
return true;
}
if (op === 'spirv.SpecConstant' || op === 'spv.SpecConstant') {
parser.parseSymbolName('sym_name', result.attributes);
if (parser.parseOptionalKeyword('spec_id')) {
parser.parseLParen();
const specId = parser.parseAttribute();
result.addAttribute('spec_id', specId);
parser.parseRParen();
}
if (parser.parseOptionalEqual()) {
const defaultValue = parser.parseAttribute();
result.addAttribute('default_value', defaultValue);
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addAttribute('type', type);
}
return true;
}
if (op === 'spirv.module' || op === 'spv.module') {
const modSymName = parser.parseOptionalSymbolName();
if (modSymName) {
result.addAttribute('sym_name', modSymName);
}
const addressingModel = parser.parseOptionalKeyword();
if (addressingModel) {
result.addAttribute('addressing_model', addressingModel);
}
const memoryModel = parser.parseOptionalKeyword();
if (memoryModel) {
result.addAttribute('memory_model', memoryModel);
}
if (parser.parseOptionalKeyword('requires')) {
const vce = parser.parseAttribute();
result.addAttribute('vce_triple', vce);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'spirv.ARM.Graph') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'spirv.ARM.GraphEntryPoint') {
const fn = parser.parseOptionalSymbolName();
result.addAttribute('fn', fn);
const interfaceVars = [];
while (parser.parseOptionalComma()) {
const varSymbol = parser.parseOptionalSymbolName();
interfaceVars.push(varSymbol);
}
result.addAttribute('interface', interfaceVars);
return true;
}
if (op === 'spirv.func' || op === 'spv.func') {
const funcSymName = parser.parseOptionalSymbolName();
if (funcSymName) {
result.addAttribute('sym_name', funcSymName);
}
const argResult = parser.parseFunctionArgumentList();
const inputs = argResult.arguments.map((a) => a.type);
const results = [];
const resultAttrs = [];
if (parser.parseOptionalArrow()) {
parser.parseFunctionResultList(results, resultAttrs);
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
const funcControl = parser.parseOptionalString();
if (funcControl !== null) {
result.addAttribute('function_control', funcControl);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
// spirv.func is IsolatedFromAbove
parser.parseOptionalRegion(result.addRegion(), undefined, /* isIsolatedNameScope */ true);
return true;
}
if (op === 'spirv.GlobalVariable' || op === 'spv.GlobalVariable') {
const gvSymName = parser.parseOptionalSymbolName();
if (gvSymName) {
result.addAttribute('sym_name', gvSymName);
}
if (parser.parseOptionalKeyword('initializer')) {
parser.parseLParen();
const initSymbol = parser.parseOptionalSymbolName();
parser.parseRParen();
result.addAttribute('initializer', initSymbol);
}
if (parser.parseOptionalKeyword('built_in')) {
parser.parseLParen();
const builtIn = parser.parseString();
parser.parseRParen();
result.addAttribute('built_in', builtIn);
}
if (parser.parseOptionalKeyword('bind')) {
parser.parseLParen();
const binding = parser.parseInteger();
parser.parseOptionalComma();
const set = parser.parseInteger();
parser.parseRParen();
result.addAttribute('descriptor_set', set);
result.addAttribute('binding', binding);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
}
return true;
}
if (op === 'spirv.EntryPoint' || op === 'spv.EntryPoint') {
const executionModel = parser.parseOptionalString();
if (executionModel !== null) {
result.addAttribute('execution_model', executionModel);
}
result.operands = [];
let epSym = parser.parseOptionalSymbolName();
while (epSym) {
result.addAttribute('fn', new _.SymbolRefAttr(epSym));
if (!parser.parseOptionalComma()) {
break;
}
epSym = parser.parseOptionalSymbolName();
}
return true;
}
if (op === 'spirv.ExecutionMode' || op === 'spv.ExecutionMode') {
const emSym = parser.parseOptionalSymbolName();
if (emSym) {
result.addAttribute('fn', new _.SymbolRefAttr(emSym));
}
const mode = parser.parseOptionalString();
if (mode !== null) {
result.addAttribute('execution_mode', mode);
}
const params = [];
while (parser.parseOptionalComma()) {
const value = parser.parseAttribute();
if (value) {
params.push(value);
} else {
break;
}
}
if (params.length > 0) {
result.addAttribute('values', params);
}
return true;
}
if (op === 'spirv.ExecutionModeId') {
const fnSym = parser.parseOptionalSymbolName();
if (fnSym) {
result.addAttribute('fn', new _.SymbolRefAttr(fnSym));
}
const mode = parser.parseOptionalString();
if (mode !== null) {
result.addAttribute('execution_mode', mode);
}
const values = [];
do {
const sym = parser.parseOptionalSymbolName();
if (!sym) {
break;
}
values.push(new _.SymbolRefAttr(sym));
} while (parser.parseOptionalComma());
if (values.length > 0) {
result.addAttribute('values', values);
}
return true;
}
if (op === 'spirv.mlir.loop' || op === 'spv.mlir.loop' || op === 'spirv.mlir.selection' || op === 'spv.mlir.selection') {
if (parser.parseOptionalKeyword('control')) {
parser.parseLParen();
const controlValue = parser.parseOptionalKeyword();
result.addAttribute('selection_control', controlValue);
parser.parseRParen();
}
result.addTypes(parser.parseOptionalArrowTypeList());
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
const arithmeticExtendedOps = new Set([
'spirv.IAddCarry', 'spv.IAddCarry',
'spirv.ISubBorrow', 'spv.ISubBorrow',
'spirv.SMulExtended', 'spv.SMulExtended',
'spirv.UMulExtended', 'spv.UMulExtended'
]);
if (arithmeticExtendedOps.has(op)) {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
const resultType = parser.parseType();
parser.resolveOperands(unresolvedOperands, [resultType, resultType], result.operands);
result.addTypes([resultType]);
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
return true;
}
if (op === 'spirv.INTEL.SubgroupBlockWrite' || op === 'spv.INTEL.SubgroupBlockWrite') {
const storageClass = parser.parseString();
result.addAttribute('storage_class', storageClass);
const ptrUnresolved = parser.parseOperand();
parser.parseComma();
const valueUnresolved = parser.parseOperand();
let ptrType = null;
let valueType = null;
if (parser.parseOptionalColon()) {
valueType = parser.parseType();
ptrType = new _.spirv.PointerType(valueType, storageClass);
}
parser.resolveOperand(ptrUnresolved, ptrType, result.operands);
parser.resolveOperand(valueUnresolved, valueType, result.operands);
return true;
}
if ((op === 'spirv.CopyMemory' || op === 'spv.CopyMemory')) {
const targetStorageClass = parser.parseString();
result.addAttribute('target_storage_class', targetStorageClass);
const targetUnresolved = parser.parseOperand();
parser.parseComma();
const sourceStorageClass = parser.parseString();
result.addAttribute('source_storage_class', sourceStorageClass);
const sourceUnresolved = parser.parseOperand();
if (parser.parseOptionalLSquare()) {
const memoryAccess = [];
if (!parser.parseOptionalRSquare()) {
do {
const str = parser.parseOptionalString();
if (str === null) {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
memoryAccess.push(String(intVal));
} else {
memoryAccess.push(str);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (memoryAccess.length > 0) {
result.addAttribute('memory_access', memoryAccess.join(', '));
}
}
if (parser.parseOptionalComma()) {
if (parser.parseOptionalLSquare()) {
const sourceMemoryAccess = [];
if (!parser.parseOptionalRSquare()) {
do {
const str = parser.parseOptionalString();
if (str === null) {
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
break;
}
sourceMemoryAccess.push(String(intVal));
} else {
sourceMemoryAccess.push(str);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (sourceMemoryAccess.length > 0) {
result.addAttribute('source_memory_access', sourceMemoryAccess.join(', '));
}
}
}
let targetType = null;
let sourceType = null;
if (parser.parseOptionalColon()) {
const elementType = parser.parseType();
targetType = new _.spirv.PointerType(elementType, targetStorageClass);
sourceType = new _.spirv.PointerType(elementType, sourceStorageClass);
}
parser.resolveOperand(targetUnresolved, targetType, result.operands);
parser.resolveOperand(sourceUnresolved, sourceType, result.operands);
return true;
}
return super.parseOperation(parser, result);
}
parseSPIRV_I32_1DArmTensor(parser, op, attrName) {
const values = [];
parser.parseCommaSeparatedList('square', () => {
values.push(parser.parseInteger());
});
op.addAttribute(attrName, values);
}
};
_.wasmssa = {};
_.wasmssa.LimitType = class extends _.Type {
constructor(min, max) {
super(null);
this.min = min;
this.max = max;
}
toString() {
const max = this.max === null ? '' : ` ${this.max}`;
return `!wasmssa.limit[${this.min}:${max}]`;
}
};
_.wasmssa.TableType = class extends _.Type {
constructor(reference, limit) {
super(null);
this.reference = reference;
this.limit = limit;
}
toString() {
return `!wasmssa.tabletype ${this.reference} ${this.limit}`;
}
};
_.WasmSSADialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'wasmssa');
this.registerCustomType('WasmSSA_LocalRef', this.parseLocalRefType.bind(this));
this.registerCustomDirective('ElseRegion', this.parseElseRegion.bind(this));
this.registerCustomAttribute('WasmSSA_ValTypeAttr', this.parseValTypeAttr.bind(this));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'local') {
parser.parseKeyword('ref');
parser.parseKeyword('to');
const elementType = parser.parseType();
return new _.Type(`!${dialect}.local ref to ${elementType}`);
}
if (mnemonic === 'funcref' || mnemonic === 'externref') {
return new _.Type(`!${dialect}.${mnemonic}`);
}
if (mnemonic === 'limit') {
return this.parseLimit(parser);
}
if (mnemonic === 'tabletype') {
const reference = parser.parseType();
const limit = this.parseLimit(parser);
return new _.wasmssa.TableType(reference, limit);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseLimit(parser) {
// assemblyFormat: `[` $min `:` $max? `]`
parser.parseLSquare();
const min = parser.parseInteger();
parser.parseColon();
const max = parser.parseOptionalInteger();
parser.parseRSquare();
return new _.wasmssa.LimitType(min, max);
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'wasmssa.import_global') {
const importName = parser.parseString();
result.addAttribute('importName', importName);
parser.parseKeyword('from');
const moduleName = parser.parseString();
result.addAttribute('moduleName', moduleName);
parser.parseKeyword('as');
parser.parseSymbolName('sym_name', result.attributes);
// An optional keyword follows: either `mutable` or a symbol
// visibility (e.g. `nested`). Only `mutable` is significant; any
// other keyword is consumed and ignored, matching the reference.
const keyword = parser.parseOptionalKeyword();
if (keyword === 'mutable') {
result.addAttribute('isMutable', new _.UnitAttr());
}
parser.parseColon();
const type = parser.parseType();
result.addAttribute('type', type);
return true;
}
if (op === 'wasmssa.global') {
if (parser.parseOptionalKeyword('exported')) {
result.addAttribute('exported', new _.UnitAttr());
}
parser.parseSymbolName('sym_name', result.attributes);
const type = parser.parseType();
result.addAttribute('type', type);
if (parser.parseOptionalKeyword('mutable')) {
result.addAttribute('isMutable', new _.UnitAttr());
}
parser.parseColon();
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'wasmssa.func') {
if (parser.parseOptionalKeyword('exported')) {
result.addAttribute('exported', new _.UnitAttr());
}
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
parseLocalRefType(parser) {
parser.parseKeyword('ref');
parser.parseKeyword('to');
const elementType = parser.parseType();
return new _.Type(`ref to ${elementType}`);
}
parseElseRegion(parser, op /*, elseName */) {
if (parser.parseOptionalKeyword('else')) {
const region = op.addRegion();
parser.parseRegion(region);
}
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'wasmssa.local') {
const typeAttr = op.attributes.get('type');
if (typeAttr) {
const elementType = typeAttr.type ? typeAttr.type : typeAttr;
op.addTypes([new _.Type(`!wasmssa<local ref to ${elementType}>`)]);
return;
}
}
if (op.name.getStringRef() === 'wasmssa.local_get' || op.name.getStringRef() === 'wasmssa.local_tee') {
const localVarEntry = vars.get('localVar');
if (localVarEntry && localVarEntry.types && localVarEntry.types.length > 0) {
const localRefType = localVarEntry.types[0];
if (localRefType) {
const typeStr = localRefType.toString();
const match = typeStr.match(/ref\s+to\s+(.+)/);
if (match) {
op.addTypes([new _.Type(match[1])]);
return;
}
}
}
}
super.inferResultTypes(op, vars);
}
parseValTypeAttr(parser) {
const type = parser.parseType();
return new _.TypeAttrOf(type);
}
};
_.CFDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'cf');
this.registerCustomDirective('SwitchOpCases', this.parseSwitchOpCases.bind(this));
}
parseSwitchOpCases(parser, op, flagTypes, defaultDestinationName, defaultOperands, defaultOperandTypes, caseValuesAttr, caseDestinationsName, caseOperands, caseOperandTypes) {
if (!parser.parseOptionalKeyword('default')) {
return false;
}
if (!parser.parseOptionalColon()) {
return false;
}
const defaultDestination = parser.parseOptionalSuccessor();
if (!defaultDestination) {
return false;
}
const defaultDest = { label: defaultDestination, arguments: [] };
if (parser.parseOptionalLParen()) {
const unresolvedDefaults = parser.parseOperandList();
for (const value of unresolvedDefaults) {
defaultDest.arguments.push({ value });
defaultOperands.push(value);
}
const defTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < defTypes.length; idx++) {
defaultOperandTypes.push(defTypes[idx]);
if (idx < defaultDest.arguments.length) {
defaultDest.arguments[idx].type = defTypes[idx];
}
}
parser.parseOptionalRParen();
}
op.successors = op.successors || [];
op.successors.push(defaultDest);
const caseValues = [];
const caseOperandSegments = [defaultDest.arguments.length];
while (parser.parseOptionalComma()) {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
caseValues.push(value);
if (!parser.parseOptionalColon()) {
break;
}
const caseDestination = parser.parseOptionalSuccessor();
if (!caseDestination) {
break;
}
const caseDest = { label: caseDestination, arguments: [] };
if (parser.parseOptionalLParen()) {
const caseOps = parser.parseOperandList();
for (const operandValue of caseOps) {
caseDest.arguments.push({ value: operandValue });
caseOperands.push(operandValue);
}
const caseTypes = parser.parseOptionalColonTypeList();
for (let idx = 0; idx < caseTypes.length; idx++) {
caseOperandTypes.push(caseTypes[idx]);
if (idx < caseDest.arguments.length) {
caseDest.arguments[idx].type = caseTypes[idx];
}
}
parser.parseOptionalRParen();
}
op.successors.push(caseDest);
caseOperandSegments.push(caseDest.arguments.length);
}
if (caseValues.length > 0) {
op.addAttribute(caseValuesAttr, caseValues);
op.addAttribute('case_operand_segments', caseOperandSegments);
}
return true;
}
};
_.pdl = {};
_.pdl.ValueType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!pdl.value';
}
};
_.pdl.TypeType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!pdl.type';
}
};
_.pdl.AttributeType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!pdl.attribute';
}
};
_.pdl.OperationType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!pdl.operation';
}
};
_.pdl.RangeType = class extends _.Type {
constructor(elementType) {
super(null);
this.elementType = elementType;
}
static getElementTypeOrSelf(type) {
if (type instanceof _.pdl.RangeType) {
return type.elementType;
}
return type;
}
toString() {
const elementStr = this.elementType?.toString() || '';
const match = elementStr.match(/^!pdl\.(.+)$/);
return `!pdl.range<${match ? match[1] : elementStr}>`;
}
};
_.pdl.PDLDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'pdl');
this.registerCustomDirective('OperationOpAttributes', this.parseOperationOpAttributes.bind(this));
this.registerCustomDirective('RangeType', this.parseRangeType.bind(this));
this.registerCustomDirective('ResultsValueType', this.parseResultsValueType.bind(this));
}
parseOperationOpAttributes(parser, op, attrOperands, attrNamesAttr) {
if (!parser.parseOptionalLBrace()) {
return true;
}
const attributeNames = [];
parser.parseCommaSeparatedList('none', () => {
const name = parser.parseAttribute();
parser.parseEqual();
const unresolvedValue = parser.parseOperand();
attrOperands.push(unresolvedValue);
attributeNames.push(name);
});
parser.parseRBrace();
if (attributeNames.length > 0) {
op.addAttribute(attrNamesAttr, attributeNames);
}
return true;
}
parseRangeType(parser, op, argumentTypes, resultTypes) {
if (argumentTypes && argumentTypes.length > 0) {
const elementType = _.pdl.RangeType.getElementTypeOrSelf(argumentTypes[0]);
resultTypes.push(new _.pdl.RangeType(elementType));
return true;
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
resultTypes.push(type);
}
return true;
}
parseResultsValueType(parser, op, indexAttr, valTypes) {
// Parses `-> type` for pdl.results operation
// If index is present, type can be !pdl.value or !pdl.range<value>
// If index is absent, type is always !pdl.range<value> (full result range)
if (!indexAttr) {
valTypes.push(new _.Type('!pdl.range<!pdl.value>'));
} else if (parser.parseOptionalArrow()) {
const type = parser.parseType();
valTypes.push(type);
}
return true;
}
};
_.PDLInterpDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'pdl_interp');
this.registerCustomDirective('CreateOperationOpAttributes', this.parseCreateOperationOpAttributes.bind(this));
this.registerCustomDirective('CreateOperationOpResults', this.parseCreateOperationOpResults.bind(this));
this.registerCustomDirective('RangeType', this.parseRangeType.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'pdl_interp.func') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'pdl_interp.foreach') {
return this.parseForeachOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseForeachOp(parser, result) {
const loopVar = parser.parseOperand();
parser.parseColon();
const loopVarType = parser.parseType();
parser.parseKeyword('in');
const range = parser.parseOperand();
parser.resolveOperand(range, null, result.operands);
{
const region = {};
parser.parseRegion(region);
if (region.blocks && region.blocks.length > 0) {
if (!region.blocks[0].arguments) {
region.blocks[0].arguments = [];
}
region.blocks[0].arguments.push({ value: loopVar, type: loopVarType });
}
result.regions.push(region);
}
if (parser.parseOptionalArrow()) {
parser.parseOptionalSuccessor();
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseCreateOperationOpAttributes(parser, op, attrOperands, attrNamesAttr) {
const attrNames = [];
if (parser.parseOptionalLBrace()) {
if (!parser.parseOptionalRBrace()) {
do {
const nameAttr = parser.parseAttribute();
parser.parseEqual();
const operand = parser.parseOperand();
attrOperands.push(operand);
attrNames.push(nameAttr);
} while (parser.parseOptionalComma());
parser.parseRBrace();
}
}
if (attrNames.length > 0) {
op.addAttribute(attrNamesAttr, attrNames);
}
}
parseCreateOperationOpResults(parser, op, resultOperands, resultTypes, inferredResultTypesAttr) {
if (!parser.parseOptionalArrow()) {
return;
}
if (parser.parseOptionalLess()) {
parser.parseKeyword('inferred');
parser.parseGreater();
op.addAttribute(inferredResultTypesAttr, true);
return;
}
parser.parseLParen();
const unresolvedOperands = parser.parseOperandList();
const types = parser.parseOptionalColonTypeList();
for (let i = 0; i < unresolvedOperands.length; i++) {
resultOperands.push(unresolvedOperands[i]);
if (i < types.length) {
resultTypes.push(types[i]);
}
}
parser.parseRParen();
}
parseRangeType(parser, op, argumentTypes, resultTypes) {
if (argumentTypes && argumentTypes.length > 0) {
const elementType = _.pdl.RangeType.getElementTypeOrSelf(argumentTypes[0]);
resultTypes.push(new _.pdl.RangeType(elementType));
return true;
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
resultTypes.push(type);
}
return true;
}
};
_.ptr = {};
_.ptr.PtrType = class extends _.Type {
constructor(memorySpace) {
super(null);
this.memorySpace = memorySpace;
}
toString() {
if (this.memorySpace === null) {
return '!ptr.ptr';
}
return `!ptr.ptr<${this.memorySpace}>`;
}
};
_.ptr.PtrDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'ptr');
this.registerCustomAttribute('EnumProp', this.parseEnumProp.bind(this));
this.registerCustomAttribute('Ptr_PtrDiffFlags', (parser, type) => this.parseEnumFlags(parser, type, '|', true));
this.registerCustomType('Ptr_PtrType', this.parsePtrType.bind(this));
}
parseEnumProp(parser, type) {
const [innerType] = type.args;
return this.parseCustomAttributeWithFallback(parser, innerType);
}
parsePtrType(parser) {
if (parser.parseOptionalLess()) {
const memorySpace = parser.parseAttribute();
parser.parseGreater();
return new _.ptr.PtrType(memorySpace);
}
return new _.ptr.PtrType(null);
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'ptr.ptr_add' && op.operands.length >= 2) {
const baseType = op.operands[0].type;
const offsetType = op.operands[1].type;
const offsetIsShaped = offsetType instanceof _.VectorType || offsetType instanceof _.RankedTensorType;
if (!offsetIsShaped) {
if (baseType) {
op.addTypes([baseType]);
}
return;
}
const baseIsShaped = baseType instanceof _.VectorType || baseType instanceof _.RankedTensorType;
if (!baseIsShaped) {
if (offsetType instanceof _.VectorType) {
op.addTypes([new _.VectorType(offsetType.dimensions, baseType, offsetType.scalableDims)]);
} else if (offsetType instanceof _.RankedTensorType) {
op.addTypes([new _.RankedTensorType(offsetType.shape, baseType, offsetType.encoding)]);
}
return;
}
if (baseType) {
op.addTypes([baseType]);
}
return;
}
super.inferResultTypes(op, vars);
}
};
_.emitc = {};
_.emitc.LValueType = class extends _.Type {
constructor(valueType) {
super(null);
this.valueType = valueType;
}
static parse(parser) {
parser.parseLess();
const valueType = parser.parseType();
parser.parseGreater();
return new _.emitc.LValueType(valueType);
}
toString() {
return `!emitc.lvalue<${this.valueType}>`;
}
};
_.emitc.EmitCDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'emitc');
this.registerCustomType('EmitC_LValueType', this.parseLValueType.bind(this));
this.registerCustomDirective('SwitchCases', this.parseSwitchCases.bind(this));
this.registerCustomDirective('EmitCGlobalOpTypeAndInitialValue', this.parseTypeAndInitialValue.bind(this));
this.registerCustomDirective('EmitCFieldOpTypeAndInitialValue', this.parseTypeAndInitialValue.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'emitc.include') {
if (parser.parseOptionalLess()) {
const include = parser.parseString();
parser.parseGreater();
result.addAttribute('is_standard_include', true);
result.addAttribute('include', include);
} else {
const include = parser.parseString();
result.addAttribute('include', include);
}
return true;
}
if (op === 'emitc.func') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'emitc.expression') {
let operand = parser.parseOptionalOperand();
while (operand) {
parser.resolveOperand(operand, null, result.operands);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
if (parser.parseOptionalKeyword('noinline')) {
result.addAttribute('do_not_inline', true);
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
result.addAttribute('type', type);
result.addTypes(type.results);
}
}
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'emitc.if') {
const thenRegion = result.addRegion();
const elseRegion = result.addRegion();
const cond = parser.parseOperand();
const i1Type = new _.IntegerType('i1');
parser.resolveOperand(cond, i1Type, result.operands);
parser.parseRegion(thenRegion);
if (parser.parseOptionalKeyword('else')) {
parser.parseRegion(elseRegion);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'emitc.do') {
const bodyRegion = result.addRegion();
const condRegion = result.addRegion();
parser.parseRegion(bodyRegion);
parser.parseKeyword('while');
parser.parseRegion(condRegion);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'emitc.for') {
const inductionVariable = parser.parseArgument();
parser.parseEqual();
const lb = parser.parseOperand();
parser.parseKeyword('to');
const ub = parser.parseOperand();
parser.parseKeyword('step');
const step = parser.parseOperand();
let type = new _.IndexType();
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
inductionVariable.type = type;
parser.resolveOperand(lb, type, result.operands);
parser.resolveOperand(ub, type, result.operands);
parser.resolveOperand(step, type, result.operands);
const region = result.addRegion();
parser.parseRegion(region, [inductionVariable]);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
return super.parseOperation(parser, result);
}
parseLValueType(parser) {
return _.emitc.LValueType.parse(parser);
}
parseSwitchCases(parser, op, casesAttr /*, caseRegionsName */) {
const caseValues = [];
while (parser.parseOptionalKeyword('case')) {
const value = parser.parseInteger();
caseValues.push(value);
const region = op.addRegion();
parser.parseRegion(region);
}
if (caseValues.length > 0) {
op.addAttribute(casesAttr, caseValues);
}
}
parseTypeAndInitialValue(parser, op, typeAttr = 'type', valueAttr = 'initial_value') {
const type = parser.parseType();
op.addAttribute(typeAttr, type);
if (parser.parseOptionalEqual()) {
const initialValue = parser.parseAttribute(type);
op.addAttribute(valueAttr, initialValue);
}
}
};
_.AsukaDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'asuka');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// https://github.com/monellz/FlashTensor/blob/main/bench/ea.mlir
// uses batch_dims and reduce_dims not valid given the assemblyFormat spec.
// Custom parsing preserves compatibility with this file.
if (op === 'asuka.dot' || op === 'asuka.add' || op === 'asuka.split' || op === 'asuka.softmax' || op === 'asuka.reduce') {
result.compatibility = true;
result.operands = parser.parseOperandList();
let attrName = parser.parseOptionalKeyword();
while (attrName) {
if (parser.parseOptionalEqual()) {
let attrValue = null;
const intVal = parser.parseOptionalInteger();
if (intVal === null) {
attrValue = parser.parseAttribute();
if (parser.parseOptionalKeyword('x')) {
const secondValue = parser.parseAttribute();
attrValue = { kind: 'pair', first: attrValue, second: secondValue };
}
} else {
attrValue = String(intVal);
}
result.addAttribute(attrName, attrValue);
parser.parseOptionalComma();
}
attrName = parser.parseOptionalKeyword();
}
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
parser.resolveOperands(result.operands, funcType.inputs);
for (const resultType of funcType.results) {
result.addTypes([resultType]);
}
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.async = {};
_.async.TokenType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!async.token';
}
};
_.async.GroupType = class extends _.Type {
constructor() {
super(null);
}
toString() {
return '!async.group';
}
};
_.async.ValueType = class extends _.Type {
constructor(valueType) {
super(null);
this.valueType = valueType;
}
toString() {
const inner = this.valueType?.toString ? this.valueType.toString() : this.valueType;
return `!async.value<${inner}>`;
}
static parse(parser) {
if (parser.parseOptionalLess()) {
const innerType = parser.parseType();
parser.parseGreater();
return new _.async.ValueType(innerType);
}
return new _.async.ValueType(null);
}
};
_.async.AsyncDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'async');
this.registerCustomDirective('AwaitResultType', this.parseAwaitResultType.bind(this));
this.registerCustomType('Async_ValueType', (parser) => _.async.ValueType.parse(parser));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'coro.handle' || mnemonic === 'coro.id' || mnemonic === 'coro.state') {
return new _.Type(`!${dialect}.${mnemonic}`);
}
if (mnemonic === 'token') {
return new _.async.TokenType();
}
if (mnemonic === 'group') {
return new _.async.GroupType();
}
if (mnemonic === 'value') {
return _.async.ValueType.parse(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'async.execute') {
return this.parseExecuteOp(parser, result);
}
if (op === 'async.func') {
return this.parseFuncOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseExecuteOp(parser, result) {
const tokenTy = new _.async.TokenType();
const tokenArgs = parser.parseOperandList('optionalSquare');
parser.resolveOperands(tokenArgs, tokenArgs.map(() => tokenTy), result.operands);
const numDependencies = tokenArgs.length;
const valueArgs = [];
const unwrappedArgs = [];
const valueTypes = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
valueArgs.push(parser.parseOperand());
parser.parseKeyword('as');
const unwrapped = parser.parseArgument();
const type = parser.parseColonType();
valueTypes.push(type);
unwrapped.type = (type instanceof _.async.ValueType && type.valueType) ? type.valueType : null;
unwrappedArgs.push(unwrapped);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.resolveOperands(valueArgs, valueTypes, result.operands);
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([numDependencies, valueArgs.length]));
const resultTypes = parser.parseOptionalArrowTypeList();
result.addTypes([tokenTy]);
result.addTypes(resultTypes);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const executeRegion = result.addRegion();
parser.parseRegion(executeRegion, unwrappedArgs);
return true;
}
parseFuncOp(parser, result) {
parser.parseOptionalVisibilityKeyword(result.attributes);
parser.parseSymbolName('sym_name', result.attributes);
const argResult = parser.parseFunctionArgumentList();
const inputs = argResult.arguments.map((a) => a.type);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const results = [];
const resultAttrs = [];
if (parser.parseOptionalArrow()) {
parser.parseFunctionResultList(results, resultAttrs);
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const funcRegion = result.addRegion();
parser.parseOptionalRegion(funcRegion);
return true;
}
parseAwaitResultType(parser, op, operandTypeArg) {
const operandType = parser.parseType();
if (Array.isArray(operandTypeArg)) {
operandTypeArg.push(operandType);
}
if (operandType instanceof _.async.ValueType && operandType.valueType) {
op.addTypes([operandType.valueType]);
}
}
};
_.ArithDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'arith');
this.registerCustomAttribute('Arith_FastMathAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
this.registerCustomAttribute('Arith_IntegerOverflowAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'arith.select') {
return this.parseSelectOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseSelectOp(parser, result) {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
let resultType = parser.parseType();
let conditionType = null;
if (parser.parseOptionalComma()) {
conditionType = resultType;
resultType = parser.parseType();
} else {
conditionType = new _.IntegerType('i1');
}
result.addTypes([resultType]);
parser.resolveOperands(unresolvedOperands, [conditionType, resultType, resultType], result.operands);
return true;
}
};
_.BuiltinDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'builtin');
this.blobManager = new Map();
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'builtin.call' || op === 'builtin.call_indirect') {
parser.parseSymbolName('callee', result.attributes);
const unresolvedOperands = parser.parseOperandList();
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
return super.parseOperation(parser, result);
}
readType(reader) {
const typeCode = reader.readVarInt();
switch (typeCode) {
case 0: { // IntegerType
const widthAndSign = reader.readVarInt();
const width = widthAndSign >> 2;
const signedness = widthAndSign & 0x3;
if (signedness === 0) {
return new _.IntegerType(`i${width}`);
}
if (signedness === 1) {
return new _.IntegerType(`si${width}`);
}
return new _.IntegerType(`ui${width}`);
}
case 1: // IndexType
return new _.IndexType();
case 2: { // FunctionType
const numInputs = reader.readVarInt();
const inputs = [];
for (let i = 0; i < numInputs; i++) {
inputs.push(reader.readType());
}
const numResults = reader.readVarInt();
const results = [];
for (let i = 0; i < numResults; i++) {
results.push(reader.readType());
}
return new _.FunctionType(inputs, results);
}
case 3: return new _.FloatType('bf16'); // BFloat16Type
case 4: return new _.FloatType('f16'); // Float16Type
case 5: return new _.FloatType('f32'); // Float32Type
case 6: return new _.FloatType('f64'); // Float64Type
case 7: return new _.FloatType('f80'); // Float80Type
case 8: return new _.FloatType('f128'); // Float128Type
case 9: { // ComplexType
const elementType = reader.readType();
return new _.ComplexType(elementType);
}
case 10: { // MemRefType
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
const layout = reader.readAttribute();
return new _.MemRefType(shape, elementType, layout, null);
}
case 11: { // MemRefTypeWithMemSpace
const memorySpace = reader.readAttribute();
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
const layout = reader.readAttribute();
return new _.MemRefType(shape, elementType, layout, memorySpace);
}
case 12: // NoneType
return new _.NoneType();
case 13: { // RankedTensorType
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
return new _.RankedTensorType(shape, elementType, null);
}
case 14: { // RankedTensorTypeWithEncoding
const encoding = reader.readAttribute();
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
return new _.RankedTensorType(shape, elementType, encoding);
}
case 15: { // TupleType
const numTypes = reader.readVarInt();
const types = [];
for (let i = 0; i < numTypes; i++) {
types.push(reader.readType());
}
return new _.TupleType(types);
}
case 16: { // UnrankedMemRefType
const elementType = reader.readType();
return new _.UnrankedMemRefType(elementType);
}
case 17: { // UnrankedMemRefTypeWithMemSpace
const memorySpace = reader.readAttribute();
const elementType = reader.readType();
return new _.UnrankedMemRefType(elementType, memorySpace);
}
case 18: { // UnrankedTensorType
const elementType = reader.readType();
return new _.UnrankedTensorType(elementType);
}
case 19: { // VectorType
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
return new _.VectorType(shape, elementType);
}
case 20: { // VectorTypeWithScalableDims
const numScalable = reader.readVarInt();
for (let i = 0; i < numScalable; i++) {
reader.readByte(); // scalableDims flags
}
const shape = reader.readSignedVarInts();
const elementType = reader.readType();
return new _.VectorType(shape, elementType);
}
default:
throw new mlir.Error(`Unsupported built-in type code '${typeCode}'.`);
}
}
readAttribute(reader) {
const typeCode = reader.readVarInt();
switch (typeCode) {
case 0: { // ArrayAttr
const count = reader.readVarInt();
const elements = [];
for (let i = 0; i < count; i++) {
elements.push(reader.readAttribute());
}
return new _.ArrayAttr(elements);
}
case 1: { // DictionaryAttr
const count = reader.readVarInt();
const attrs = new Map();
for (let i = 0; i < count; i++) {
const nameAttr = reader.readAttribute();
const valueAttr = reader.readAttribute();
const name = nameAttr && nameAttr.value ? nameAttr.value : `attr_${i}`;
attrs.set(name, valueAttr);
}
return { name: 'dictionary', value: attrs };
}
case 2: { // StringAttr
const value = reader.readString();
return new _.StringAttr(value);
}
case 3: { // StringAttrWithType
const value = reader.readString();
const type = reader.readType();
return new _.StringAttr(value, type);
}
case 4: { // FlatSymbolRefAttr
const value = reader.readString();
return new _.SymbolRefAttr(`@${value}`);
}
case 5: { // SymbolRefAttr
const root = reader.readString();
const numNested = reader.readVarInt();
let value = `@${root}`;
for (let i = 0; i < numNested; i++) {
value += `::@${reader.readString()}`;
}
return new _.SymbolRefAttr(value);
}
case 6: { // TypeAttr
const type = reader.readType();
return new _.TypeAttrOf(type);
}
case 7: // UnitAttr
return new _.UnitAttr();
case 8: { // IntegerAttr
const getIntegerBitWidth = (type) => {
const str = type ? type.toString() : '';
const match = str.match(/^[su]?i(\d+)$/);
if (match) {
return parseInt(match[1], 10);
}
if (str === 'index') {
return 64;
}
throw new mlir.Error(`Unsupported integer type '${str}'.`);
};
const type = reader.readType();
const bitWidth = getIntegerBitWidth(type);
let value = null;
if (bitWidth <= 8) {
value = BigInt(reader.readByte());
} else if (bitWidth <= 64) {
value = reader.readSignedVarInt();
} else {
const numWords = reader.readVarInt();
value = 0n;
for (let i = 0; i < numWords; i++) {
const word = reader.readSignedVarInt();
value |= (word << BigInt(i * 64));
}
}
return new _.IntegerAttr(type, value);
}
case 9: { // FloatAttr
const type = reader.readType();
const value = reader.readAPFloatWithKnownSemantics(type);
return new _.FloatAttr(type, value);
}
case 10: { // CallSiteLoc
const caller = reader.readAttribute();
const callee = reader.readAttribute();
return new _.CallSiteLoc(callee, caller);
}
case 11: { // FileLineColLoc
const filename = reader.readString();
const line = reader.readVarInt();
const col = reader.readVarInt();
return new _.FileLineColLoc(null, filename, line, col);
}
case 12: { // FusedLoc
const count = reader.readVarInt();
const locations = [];
for (let i = 0; i < count; i++) {
locations.push(reader.readAttribute());
}
return new _.FusedLoc(locations, null);
}
case 13: { // FusedLocWithMetadata
const metadata = reader.readAttribute();
const count = reader.readVarInt();
const locations = [];
for (let i = 0; i < count; i++) {
locations.push(reader.readAttribute());
}
return new _.FusedLoc(locations, metadata);
}
case 14: { // NameLoc
const nameAttr = reader.readAttribute();
const childLoc = reader.readAttribute();
const nameStr = nameAttr && nameAttr.value !== undefined ? nameAttr.value : '';
return new _.NameLoc(nameStr, childLoc);
}
case 15: // UnknownLoc
return new _.UnknownLoc();
case 16: { // DenseResourceElementsAttr
const type = reader.readType();
const resource = reader.readResourceHandle();
const blobData = resource && resource.value && resource.value.kind === 'blob' ? resource.value.data : null;
const handle = new _.DenseResourceElementsHandle(resource ? resource.key : 'unknown', blobData);
return new _.DenseResourceElementsAttr(type, handle);
}
case 17: { // DenseArrayAttr
const type = reader.readType();
const size = reader.readVarInt();
const blob = reader.readBlob();
return new _.DenseArrayAttr(type, size, blob);
}
case 18: { // DenseIntOrFPElementsAttr
const type = reader.readType();
const blob = reader.readBlob();
return new _.DenseElementsAttr(blob, type);
}
case 19: { // DenseStringElementsAttr
const type = reader.readType();
const isSplat = reader.readVarInt() !== 0;
const count = reader.readVarInt();
const strings = [];
for (let i = 0; i < count; i++) {
strings.push(reader.readString());
}
return { name: 'dense_string', value: strings, type, isSplat };
}
case 20: { // SparseElementsAttr
const type = reader.readType();
const indices = reader.readAttribute();
const values = reader.readAttribute();
return new _.SparseElementsAttr(type, indices, values);
}
case 21: { // DistinctAttr
const referencedAttr = reader.readAttribute();
return { name: 'distinct', value: referencedAttr };
}
case 22: { // FileLineColRange
const filename = reader.readString();
const numLocs = reader.readVarInt();
const locs = [];
for (let i = 0; i < numLocs; i++) {
locs.push(reader.readVarInt());
}
if (locs.length >= 4) {
return new _.FileLineColRange(filename, locs[0], locs[1], locs[2], locs[3]);
}
if (locs.length === 2) {
return new _.FileLineColLoc(null, filename, locs[0], locs[1]);
}
return new _.FileLineColLoc(null, filename, locs[0] || 0, 0);
}
default:
return { name: 'builtin', value: `<builtin code ${typeCode}>` };
}
}
declareResource(key) {
if (!this.blobManager.has(key)) {
this.blobManager.set(key, new _.DenseResourceElementsHandle(key));
}
return this.blobManager.get(key);
}
getResourceKey(handle) {
return handle.key;
}
parseResource(entry) {
const blob = entry.parseAsBlob();
this.blobManager.get(entry.key).blob = blob;
}
};
_.BufferizationDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'bufferization');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'bufferization.alloc_tensor') {
if (!parser.parseOptionalLParen()) {
return false;
}
const unresolvedDynamicDims = [];
if (!parser.parseOptionalRParen()) {
let operand = parser.parseOptionalOperand();
while (operand) {
unresolvedDynamicDims.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.parseRParen();
}
let unresolvedCopy = null;
if (parser.parseOptionalKeyword('copy')) {
parser.parseLParen();
unresolvedCopy = parser.parseOperand();
parser.parseRParen();
}
let unresolvedSizeHint = null;
if (parser.parseOptionalKeyword('size_hint')) {
parser.parseEqual();
unresolvedSizeHint = parser.parseOperand();
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const resultType = parser.parseType();
const indexType = new _.IndexType();
parser.resolveOperands(unresolvedDynamicDims, unresolvedDynamicDims.map(() => indexType), result.operands);
if (unresolvedCopy) {
parser.resolveOperand(unresolvedCopy, resultType, result.operands);
}
if (unresolvedSizeHint) {
parser.resolveOperand(unresolvedSizeHint, indexType, result.operands);
}
if (result.types.length === 0) {
result.addTypes([resultType]);
} else {
result.types[0] = resultType;
}
}
return true;
}
// bufferization.to_memref %tensor read_only : tensor_type to memref_type
if (op === 'bufferization.to_memref') {
const unresolvedOperand = parser.parseOptionalOperand();
if (parser.parseOptionalKeyword('read_only')) {
result.addAttribute('read_only', true);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const sourceType = parser.parseType();
result.addAttribute('source_type', sourceType);
if (unresolvedOperand) {
parser.resolveOperand(unresolvedOperand, sourceType, result.operands);
}
parser.parseKeyword('to');
const destType = parser.parseType();
result.addTypes([destType]);
} else if (unresolvedOperand) {
parser.resolveOperand(unresolvedOperand, null, result.operands);
}
return true;
}
return super.parseOperation(parser, result);
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'bufferization.dealloc') {
// DeallocOp::inferReturnTypes - one i1 per retained memref
const retainedEntry = vars.get('retained');
const numRetained = retainedEntry?.operands?.length || 0;
const i1Type = new _.IntegerType('i1');
for (let i = 0; i < numRetained; i++) {
op.addTypes([i1Type]);
}
return;
}
super.inferResultTypes(op, vars);
}
};
_.SCFDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'scf');
this.registerCustomDirective('SwitchCases', this.parseSwitchCases.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'scf.for') {
return this.parseForOp(parser, result);
}
if (op === 'scf.if') {
return this.parseIfOp(parser, result);
}
if (op === 'scf.while') {
return this.parseWhileOp(parser, result);
}
if (op === 'scf.forall') {
return this.parseForallOp(parser, result);
}
if (op === 'scf.forall.in_parallel') {
return this.parseInParallelOp(parser, result);
}
if (op === 'scf.parallel') {
return this.parseParallelOp(parser, result);
}
if (op === 'scf.execute_region') {
return this.parseExecuteRegionOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseForOp(parser, result) {
if (parser.parseOptionalKeyword('unsigned')) {
result.addAttribute('unsignedCmp', true);
}
const inductionVariable = parser.parseArgument();
parser.parseEqual();
const lb = parser.parseOperand();
parser.parseKeyword('to');
const ub = parser.parseOperand();
parser.parseKeyword('step');
const step = parser.parseOperand();
const regionArgs = [inductionVariable];
const operands = [];
const hasIterArgs = parser.parseOptionalKeyword('iter_args');
if (hasIterArgs) {
parser.parseAssignmentList(regionArgs, operands);
result.addTypes(parser.parseArrowTypeList());
}
const type = parser.parseOptionalColon() ? parser.parseType() : new _.IndexType();
inductionVariable.type = type;
for (let i = 0; i < operands.length; i++) {
regionArgs[i + 1].type = result.types[i];
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
parser.resolveOperands([lb, ub, step], [type, type, type], result.operands);
if (hasIterArgs) {
parser.resolveOperands(operands, result.types, result.operands);
}
parser.parseOptionalAttrDict(result.attributes);
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([1, 1, 1, operands.length]));
return true;
}
parseIfOp(parser, result) {
const thenRegion = result.addRegion();
const elseRegion = result.addRegion();
const unresolvedCond = parser.parseOptionalOperand();
if (!unresolvedCond) {
return false;
}
const i1Type = new _.IntegerType('i1');
parser.resolveOperands([unresolvedCond], [i1Type], result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
parser.parseRegion(thenRegion);
if (parser.parseOptionalKeyword('else')) {
parser.parseRegion(elseRegion);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseWhileOp(parser, result) {
const regionArgs = [];
const operands = [];
parser.parseOptionalAssignmentList(regionArgs, operands);
parser.parseColon();
const functionType = parser.parseType();
result.addTypes(functionType.results);
parser.resolveOperands(operands, functionType.inputs, result.operands);
for (let i = 0; i < regionArgs.length; i++) {
regionArgs[i].type = functionType.inputs[i];
}
const before = result.addRegion();
parser.parseRegion(before, regionArgs);
parser.parseKeyword('do');
const after = result.addRegion();
parser.parseRegion(after);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
parseForallOp(parser, result) {
const indexType = new _.IndexType();
const ivs = parser.parseArgumentList('paren');
const parseBoundsList = () => {
const dynamicBounds = [];
const staticBounds = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOptionalOperand();
if (operand) {
dynamicBounds.push(operand);
staticBounds.push(_.ShapedType.kDynamic);
} else {
staticBounds.push(parser.parseInteger('int64'));
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
return { dynamicBounds, staticBounds };
};
let dynamicLbs = [];
let dynamicSteps = [];
let dynamicUbs = [];
let staticLbs = null;
let staticSteps = null;
let staticUbs = null;
if (parser.parseOptionalKeyword('in')) {
const ubs = parseBoundsList();
dynamicUbs = ubs.dynamicBounds;
staticUbs = ubs.staticBounds;
parser.resolveOperands(dynamicUbs, dynamicUbs.map(() => indexType), result.operands);
const numLoops = ivs.length;
staticLbs = new Array(numLoops).fill(0n);
staticSteps = new Array(numLoops).fill(1n);
} else {
parser.parseEqual();
const lbs = parseBoundsList();
dynamicLbs = lbs.dynamicBounds;
staticLbs = lbs.staticBounds;
parser.resolveOperands(dynamicLbs, dynamicLbs.map(() => indexType), result.operands);
parser.parseKeyword('to');
const ubs = parseBoundsList();
dynamicUbs = ubs.dynamicBounds;
staticUbs = ubs.staticBounds;
parser.resolveOperands(dynamicUbs, dynamicUbs.map(() => indexType), result.operands);
parser.parseKeyword('step');
const steps = parseBoundsList();
dynamicSteps = steps.dynamicBounds;
staticSteps = steps.staticBounds;
parser.resolveOperands(dynamicSteps, dynamicSteps.map(() => indexType), result.operands);
}
const regionOutArgs = [];
const outOperands = [];
if (parser.parseOptionalKeyword('shared_outs')) {
parser.parseAssignmentList(regionOutArgs, outOperands);
result.addTypes(parser.parseOptionalArrowTypeList());
parser.resolveOperands(outOperands, result.types, result.operands);
}
const regionArgs = [];
for (const iv of ivs) {
iv.type = indexType;
regionArgs.push(iv);
}
for (let i = 0; i < regionOutArgs.length; i++) {
regionOutArgs[i].type = result.types[i];
regionArgs.push(regionOutArgs[i]);
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
parser.parseOptionalAttrDict(result.attributes);
result.addAttribute('staticLowerBound', staticLbs);
result.addAttribute('staticUpperBound', staticUbs);
result.addAttribute('staticStep', staticSteps);
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([dynamicLbs.length, dynamicUbs.length, dynamicSteps.length, outOperands.length]));
return true;
}
parseParallelOp(parser, result) {
const indexType = new _.IndexType();
const ivs = parser.parseArgumentList('paren');
parser.parseEqual();
const lower = parser.parseOperandList('paren');
parser.resolveOperands(lower, lower.map(() => indexType), result.operands);
parser.parseKeyword('to');
const upper = parser.parseOperandList('paren');
parser.resolveOperands(upper, upper.map(() => indexType), result.operands);
parser.parseKeyword('step');
const steps = parser.parseOperandList('paren');
parser.resolveOperands(steps, steps.map(() => indexType), result.operands);
const initVals = [];
if (parser.parseOptionalKeyword('init')) {
initVals.push(...parser.parseOperandList('paren'));
}
result.addTypes(parser.parseOptionalArrowTypeList());
for (const iv of ivs) {
iv.type = indexType;
}
const region = result.addRegion();
parser.parseRegion(region, ivs);
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr([lower.length, upper.length, steps.length, initVals.length]));
parser.parseOptionalAttrDict(result.attributes);
parser.resolveOperands(initVals, result.types, result.operands);
return true;
}
parseInParallelOp(parser, result) {
// scf.forall.in_parallel { region }
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseSwitchCases(parser, op, casesAttrName) {
const caseValues = [];
while (parser.parseOptionalKeyword('case')) {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
caseValues.push(value);
const region = op.addRegion();
parser.parseRegion(region);
}
if (casesAttrName) {
op.addAttribute(casesAttrName, caseValues);
}
}
parseExecuteRegionOp(parser, result) {
result.addTypes(parser.parseOptionalArrowTypeList());
if (parser.parseOptionalKeyword('no_inline')) {
result.addAttribute('no_inline', true);
}
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
};
_.ShapeDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'shape');
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
const simpleTypes = ['shape', 'witness', 'size', 'value_shape'];
if (simpleTypes.includes(mnemonic)) {
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'shape.func') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'shape.assuming') {
return this.parseAssumingOp(parser, result);
}
if (op === 'shape.const_shape') {
return this.parseConstShapeOp(parser, result);
}
if (op === 'shape.reduce') {
return this.parseReduceOp(parser, result);
}
if (op === 'shape.function_library') {
return this.parseFunctionLibraryOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseAssumingOp(parser, result) {
const unresolvedWitness = parser.parseOptionalOperand();
if (!unresolvedWitness) {
return false;
}
const witnessType = new _.Type('!shape.witness');
parser.resolveOperand(unresolvedWitness, witnessType, result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseConstShapeOp(parser, result) {
parser.parseOptionalAttrDict(result.attributes);
const extents = parser.parseAttribute();
result.addAttribute('shape', extents);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
}
return true;
}
parseReduceOp(parser, result) {
if (!parser.parseOptionalLParen()) {
return false;
}
const unresolvedOperands = [];
let operand = parser.parseOptionalOperand();
while (operand) {
unresolvedOperands.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.parseOptionalRParen();
let shapeType = new _.Type('!shape.shape');
if (parser.parseOptionalColon()) {
shapeType = parser.parseType();
}
const resultTypes = [];
{
const types = parser.parseOptionalArrowTypeList();
result.addTypes(types);
resultTypes.push(...types);
}
// First operand is the shape, rest are init values with result types
if (unresolvedOperands.length > 0) {
parser.resolveOperand(unresolvedOperands[0], shapeType, result.operands);
for (let i = 1; i < unresolvedOperands.length; i++) {
const initType = resultTypes[i - 1] || null;
parser.resolveOperand(unresolvedOperands[i], initType, result.operands);
}
}
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseFunctionLibraryOp(parser, result) {
parser.parseSymbolName('sym_name', result.attributes);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
if (parser.parseOptionalKeyword('mapping')) {
const mapping = parser.parseAttribute();
result.addAttribute('mapping', mapping);
}
return true;
}
};
_.SparseTensorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'sparse_tensor');
this.registerCustomDirective('LevelRange', this.parseLevelRange.bind(this));
}
parseLevelRange(parser, op, startAttr, endAttr) {
const loLvl = parser.parseInteger();
const hiLvl = parser.parseOptionalKeyword('to') ? parser.parseInteger() : loLvl + 1;
if (startAttr && endAttr) {
op.addAttribute(startAttr, loLvl);
op.addAttribute(endAttr, hiLvl);
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'sparse_tensor.iterate') {
return this.parseIterateOp(parser, result);
}
if (op === 'sparse_tensor.coiterate') {
return this.parseCoIterateOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseIterateOp(parser, result) {
const iteratorArg = parser.parseOptionalOperand();
if (!iteratorArg) {
return false;
}
const regionArgs = [];
regionArgs.push({ name: iteratorArg.name, type: null }); // type determined by tensor
if (!parser.parseOptionalKeyword('in')) {
return false;
}
const unresolvedTensor = parser.parseOptionalOperand();
if (!unresolvedTensor) {
return false;
}
const iterArgNames = [];
const initValues = [];
if (parser.parseOptionalKeyword('at')) {
parser.parseOptionalLParen();
do {
const _atArg = parser.parseOptionalOperand();
if (!_atArg && !parser.parseOptionalKeyword()) {
break;
}
} while (parser.parseOptionalComma());
parser.parseOptionalRParen();
}
if (parser.parseOptionalKeyword('iter_args')) {
parser.parseOptionalLParen();
let iterArg = parser.parseOptionalOperand();
while (iterArg) {
iterArgNames.push(iterArg.name);
if (parser.parseOptionalEqual()) {
initValues.push(parser.parseOperand());
}
if (!parser.parseOptionalComma()) {
break;
}
iterArg = parser.parseOptionalOperand();
}
parser.parseOptionalRParen();
}
let tensorType = null;
if (parser.parseOptionalColon()) {
tensorType = parser.parseType();
}
const resultTypes = [];
{
const types = parser.parseOptionalArrowTypeList();
result.addTypes(types);
resultTypes.push(...types);
}
// Add iter_args to region args with their result types
for (let i = 0; i < iterArgNames.length; i++) {
const argType = resultTypes[i] || null;
regionArgs.push({ name: iterArgNames[i], type: argType });
}
parser.resolveOperand(unresolvedTensor, tensorType, result.operands);
// iter_args block args don't go to operands, but init values do
for (let i = 0; i < initValues.length; i++) {
const initType = resultTypes[i] || null;
parser.resolveOperand(initValues[i], initType, result.operands);
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseCoIterateOp(parser, result) {
if (!parser.parseOptionalLParen()) {
return false;
}
const unresolvedTensors = [];
let _tensor = parser.parseOptionalOperand();
while (_tensor) {
unresolvedTensors.push(_tensor);
if (!parser.parseOptionalComma()) {
break;
}
_tensor = parser.parseOptionalOperand();
}
parser.parseOptionalRParen();
if (parser.parseOptionalKeyword('at')) {
parser.parseOptionalLParen();
do {
const _coAtArg = parser.parseOptionalOperand();
if (!_coAtArg && !parser.parseOptionalKeyword()) {
break;
}
} while (parser.parseOptionalComma());
parser.parseOptionalRParen();
}
const iterArgNames = [];
const initValues = [];
if (parser.parseOptionalKeyword('iter_args')) {
parser.parseOptionalLParen();
let coIterArg = parser.parseOptionalOperand();
while (coIterArg) {
iterArgNames.push(coIterArg.name);
if (parser.parseOptionalEqual()) {
initValues.push(parser.parseOperand());
}
if (!parser.parseOptionalComma()) {
break;
}
coIterArg = parser.parseOptionalOperand();
}
parser.parseOptionalRParen();
}
const tensorTypes = [];
if (parser.parseOptionalColon()) {
tensorTypes.push(...parser.parseCommaSeparatedList('optionalParen', () => parser.parseType()));
}
const resultTypes = [];
{
const types = parser.parseOptionalArrowTypeList();
result.addTypes(types);
resultTypes.push(...types);
}
const regionArgs = [];
for (let i = 0; i < iterArgNames.length; i++) {
const argType = resultTypes[i] || null;
regionArgs.push({ name: iterArgNames[i], type: argType });
}
for (let i = 0; i < unresolvedTensors.length; i++) {
const tensorType = tensorTypes[i] || null;
parser.resolveOperand(unresolvedTensors[i], tensorType, result.operands);
}
for (let i = 0; i < initValues.length; i++) {
const initType = resultTypes[i] || null;
parser.resolveOperand(initValues[i], initType, result.operands);
}
while (parser.parseOptionalKeyword('case')) {
const caseArgs = [...regionArgs]; // Start with iter_args
do {
const _caseOp = parser.parseOptionalOperand();
if (_caseOp) {
caseArgs.push({ name: _caseOp.name, type: null });
} else if (!parser.parseOptionalKeyword()) {
break;
}
} while (parser.parseOptionalComma());
const region = result.addRegion();
parser.parseRegion(region, caseArgs);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
};
_.FuncDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'func');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'func.func') {
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
};
_.GpuDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'gpu');
this.registerCustomDirective('AllReduceOperation', this.parseAllReduceOperation.bind(this));
this.registerCustomDirective('LaunchFuncOperands', this.parseLaunchFuncOperands.bind(this));
this.registerCustomDirective('AsyncDependencies', this.parseAsyncDependencies.bind(this));
this.registerCustomDirective('LaunchDimType', this.parseLaunchDimType.bind(this));
this.registerCustomDirective('OffloadingHandler', this.parseOffloadingHandler.bind(this));
}
parseAllReduceOperation(parser, op, attrName = 'op') {
const validOps = ['add', 'mul', 'minui', 'minsi', 'minnumf', 'maxui', 'maxsi', 'maxnumf', 'and', 'or', 'xor', 'minimumf', 'maximumf'];
const opName = parser.parseOptionalKeyword();
if (opName && validOps.includes(opName)) {
op.addAttribute(attrName, opName);
}
}
parseLaunchDimType(parser, op, typeArg1, typeArg2, clusterTypeArg1, clusterTypeArg2, clusterTypeArg3) {
// typeArg1 = type($gridSizeX), typeArg2 = ref($clusterSizeX)
// clusterTypeArg1/2/3 = type($clusterSizeX/Y/Z)
let dimType = new _.IndexType();
if (parser.parseOptionalColon()) {
dimType = parser.parseType();
}
// Push type to gridSizeX types array
if (Array.isArray(typeArg1)) {
typeArg1.push(dimType);
}
// If clusters are present (ref($clusterSizeX) has values), push to cluster type arrays
const hasCluster = Array.isArray(typeArg2) && typeArg2.length > 0;
if (hasCluster) {
if (Array.isArray(clusterTypeArg1)) {
clusterTypeArg1.push(dimType);
}
if (Array.isArray(clusterTypeArg2)) {
clusterTypeArg2.push(dimType);
}
if (Array.isArray(clusterTypeArg3)) {
clusterTypeArg3.push(dimType);
}
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'gpu.func') {
parser.parseOptionalVisibilityKeyword(result.attributes);
parser.parseSymbolName('sym_name', result.attributes);
const sig = parser.parseFunctionSignatureWithArguments(false);
const argTypes = sig.arguments.map((a) => a.type);
const type = new _.FunctionType(argTypes, sig.resultTypes);
result.addAttribute('function_type', new _.TypeAttrOf(type));
const allArgs = [...sig.arguments];
if (parser.parseOptionalKeyword('workgroup')) {
const workgroupResult = parser.parseFunctionArgumentList(false);
allArgs.push(...workgroupResult.arguments);
}
if (parser.parseOptionalKeyword('private')) {
const privateResult = parser.parseFunctionArgumentList(false);
allArgs.push(...privateResult.arguments);
}
if (parser.parseOptionalKeyword('kernel')) {
result.addAttribute('gpu.kernel', true);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
// gpu.func is IsolatedFromAbove
parser.parseRegion(region, allArgs, /* isIsolatedNameScope */ true);
return true;
}
if (op === 'gpu.launch') {
const indexType = new _.IndexType();
if (parser.parseOptionalKeyword('async')) {
if (parser.getNumResults() === 0) {
parser.emitError(`Operation '${op}' needs to be named when marked 'async'`);
}
result.addTypes([new _.Type('!gpu.async.token')]);
}
const asyncTokenType = new _.Type('!gpu.async.token');
const asyncDeps = parser.parseOperandList('optionalSquare');
parser.resolveOperands(asyncDeps, asyncDeps.map(() => asyncTokenType), result.operands);
let hasCluster = false;
if (parser.parseOptionalKeyword('clusters')) {
hasCluster = true;
this.parseSizeAssignment(parser, result, indexType);
parser.parseKeyword('in');
this.parseSizeAssignment(parser, result, indexType);
}
parser.parseKeyword('blocks');
this.parseSizeAssignment(parser, result, indexType);
parser.parseKeyword('in');
this.parseSizeAssignment(parser, result, indexType);
parser.parseKeyword('threads');
this.parseSizeAssignment(parser, result, indexType);
parser.parseKeyword('in');
this.parseSizeAssignment(parser, result, indexType);
let hasDynamicSharedMemorySize = false;
if (parser.parseOptionalKeyword('dynamic_shared_memory_size')) {
hasDynamicSharedMemorySize = true;
const operand = parser.parseOperand();
parser.resolveOperand(operand, new _.IntegerType('i32'), result.operands);
}
if (parser.parseOptionalKeyword('module')) {
parser.parseLParen();
const moduleSymbol = parser.parseOptionalSymbolName();
result.addAttribute('module', moduleSymbol);
parser.parseRParen();
}
if (parser.parseOptionalKeyword('function')) {
parser.parseLParen();
const funcSymbol = parser.parseOptionalSymbolName();
result.addAttribute('function', funcSymbol);
parser.parseRParen();
}
const regionArgs = [];
let numWorkgroupAttrs = 0;
if (parser.parseOptionalKeyword('workgroup')) {
const workgroupArgs = parser.parseArgumentList('paren', true);
numWorkgroupAttrs = workgroupArgs.length;
regionArgs.push(...workgroupArgs);
}
result.addAttribute('workgroup_attributions', numWorkgroupAttrs);
if (parser.parseOptionalKeyword('private')) {
const privateArgs = parser.parseArgumentList('paren', true);
regionArgs.push(...privateArgs);
}
if (parser.parseOptionalKeyword('cooperative')) {
result.addAttribute('cooperative', new _.UnitAttr());
}
const region = result.addRegion();
parser.parseRegion(region, regionArgs.length > 0 ? regionArgs : undefined, /* isIsolatedNameScope */ true);
parser.parseOptionalAttrDict(result.attributes);
const segmentSizes = [asyncDeps.length, 1, 1, 1, 1, 1, 1, hasCluster ? 1 : 0, hasCluster ? 1 : 0, hasCluster ? 1 : 0, hasDynamicSharedMemorySize ? 1 : 0];
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr(segmentSizes));
return true;
}
if (op === 'gpu.warp_execute_on_lane_0') {
return this.parseWarpExecuteOnLane0Op(parser, result);
}
return super.parseOperation(parser, result);
}
parseWarpExecuteOnLane0Op(parser, result) {
parser.parseLParen();
const unresolvedLaneId = parser.parseOperand();
const indexType = new _.IndexType();
parser.resolveOperand(unresolvedLaneId, indexType, result.operands);
parser.parseRParen();
parser.parseLSquare();
const warpSize = parser.parseInteger();
result.addAttribute('warp_size', warpSize);
parser.parseRSquare();
if (parser.parseOptionalKeyword('args')) {
parser.parseLParen();
const unresolvedArgs = parser.parseOperandList('none');
if (parser.parseOptionalColon()) {
const types = parser.parseTypeList();
parser.resolveOperands(unresolvedArgs, types, result.operands);
} else {
for (const arg of unresolvedArgs) {
parser.resolveOperand(arg, null, result.operands);
}
}
parser.parseRParen();
}
const arrowTypes = parser.parseOptionalArrowTypeList();
if (arrowTypes.length > 0) {
if (result.types.length > 0) {
result.addTypes(arrowTypes);
} else {
for (const type of arrowTypes) {
result.addTypes([type]);
}
}
}
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseSizeAssignment(parser, op, indexType) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const blockArg = parser.parseOptionalOperand();
if (!blockArg) {
break;
}
if (parser.parseOptionalEqual()) {
const operand = parser.parseOperand();
parser.resolveOperand(operand, indexType, op.operands);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parseLaunchFuncOperands(parser, op /*, args */) {
if (parser.parseOptionalKeyword('args')) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(operand, type, op.operands);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
}
parseAsyncDependencies(parser, op, asyncTokenTypes, asyncDependencies) {
// custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
// If 'async' keyword is present, operation has async result token
const hasAsync = parser.parseOptionalKeyword('async');
if (hasAsync && Array.isArray(asyncTokenTypes)) {
asyncTokenTypes.push(new _.Type('!gpu.async.token'));
}
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
if (Array.isArray(asyncDependencies)) {
asyncDependencies.push(parser.parseOperand());
} else {
parser.parseOperand();
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
}
parseOffloadingHandler(parser /*, op, args */) {
if (parser.parseOptionalLess()) {
parser.parseAttribute();
parser.parseGreater();
}
}
};
_.ArmSMEDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'arm_sme');
this.registerCustomAttribute('ArmSME_TypeSizeAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
this.registerCustomAttribute('ArmSME_TileSliceLayoutAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
this.registerCustomAttribute('ArmSME_CombiningKindAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'arm_sme.outerproduct' && op.operands.length >= 2) {
const lhsType = op.operands[0].type;
if (lhsType instanceof _.VectorType && lhsType.shape.length === 1 && lhsType.scalableDims?.[0]) {
const size = lhsType.shape[0];
const tileType = new _.VectorType([size, size], lhsType.elementType, [true, true]);
op.addTypes([tileType]);
return;
}
}
super.inferResultTypes(op, vars);
}
};
_.ArmNeonDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'arm_neon');
}
};
_.ArmSVEDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'arm_sve');
}
};
_.AMDGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'amdgpu');
this.registerCustomDirective('MNKDimensionList', this.parseMNKDimensionList.bind(this));
this.registerCustomAttribute('ROCDL_SchedGroupMaskAttr', this.parseEnumFlagsPipe.bind(this));
}
parseMNKDimensionList(parser, op, mAttr, nAttr, kAttr) {
const dimInfo = parser.parseDimensionList(false, false);
const dims = dimInfo.dimensions;
if (dims.length >= 3 && mAttr && nAttr && kAttr) {
op.addAttribute(mAttr, dims[0]);
op.addAttribute(nAttr, dims[1]);
op.addAttribute(kAttr, dims[2]);
}
}
};
_.nvgpu = {};
_.nvgpu.WarpgroupAccumulatorType = class extends _.Type {
constructor(fragmented) {
super(null);
this.fragmented = fragmented;
}
static parse(parser) {
parser.parseLess();
parser.parseKeyword('fragmented');
parser.parseEqual();
const fragmented = parser.parseType();
parser.parseGreater();
return new _.nvgpu.WarpgroupAccumulatorType(fragmented);
}
toString() {
return `!nvgpu.warpgroup.accumulator<fragmented = ${this.fragmented}>`;
}
};
_.nvgpu.WarpgroupMatrixDescriptorType = class extends _.Type {
constructor(tensor) {
super(null);
this.tensor = tensor;
}
static parse(parser) {
parser.parseLess();
parser.parseKeyword('tensor');
parser.parseEqual();
const tensor = parser.parseType();
parser.parseGreater();
return new _.nvgpu.WarpgroupMatrixDescriptorType(tensor);
}
toString() {
return `!nvgpu.warpgroup.descriptor<tensor = ${this.tensor}>`;
}
};
_.nvgpu.MBarrierGroupType = class extends _.Type {
constructor(memorySpace, numBarriers) {
super(null);
this.memorySpace = memorySpace;
this.numBarriers = numBarriers;
}
static parse(parser) {
parser.parseLess();
let memorySpace = null;
let numBarriers = 1;
do {
const key = parser.parseKeyword();
parser.parseEqual();
if (key === 'memorySpace') {
memorySpace = parser.parseAttribute();
} else if (key === 'num_barriers') {
numBarriers = parser.parseInteger();
}
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.nvgpu.MBarrierGroupType(memorySpace, numBarriers);
}
toString() {
const parts = [];
if (this.memorySpace !== null) {
parts.push(`memorySpace = ${this.memorySpace}`);
}
if (this.numBarriers !== 1) {
parts.push(`num_barriers = ${this.numBarriers}`);
}
return `!nvgpu.mbarrier.group<${parts.join(', ')}>`;
}
};
_.nvgpu.TensorMapDescriptorType = class extends _.Type {
constructor(tensor, swizzle, l2promo, oob, interleave) {
super(null);
this.tensor = tensor;
this.swizzle = swizzle;
this.l2promo = l2promo;
this.oob = oob;
this.interleave = interleave;
}
static parse(parser) {
parser.parseLess();
let tensor = null;
let swizzle = null;
let l2promo = null;
let oob = null;
let interleave = null;
do {
const key = parser.parseKeyword();
parser.parseEqual();
switch (key) {
case 'tensor': tensor = parser.parseType(); break;
case 'swizzle': swizzle = parser.parseKeyword(); break;
case 'l2promo': l2promo = parser.parseKeyword(); break;
case 'oob': oob = parser.parseKeyword(); break;
case 'interleave': interleave = parser.parseKeyword(); break;
default: break;
}
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.nvgpu.TensorMapDescriptorType(tensor, swizzle, l2promo, oob, interleave);
}
toString() {
return `!nvgpu.tensormap.descriptor<tensor = ${this.tensor}, swizzle = ${this.swizzle}, l2promo = ${this.l2promo}, oob = ${this.oob}, interleave = ${this.interleave}>`;
}
};
_.nvgpu.NVGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'nvgpu');
this.registerCustomType('NVGPU_TensorMapDescriptor', (parser) => _.nvgpu.TensorMapDescriptorType.parse(parser));
this.registerCustomType('NVGPU_WarpgroupAccumulator', this.parseWarpgroupAccumulator.bind(this));
this.registerCustomType('NVGPU_WarpgroupMatrixDescriptor', this.parseWarpgroupMatrixDescriptor.bind(this));
this.registerCustomType('NVGPU_MBarrierGroup', this.parseMBarrierGroup.bind(this));
}
parseWarpgroupAccumulator(parser) {
return _.nvgpu.WarpgroupAccumulatorType.parse(parser);
}
parseWarpgroupMatrixDescriptor(parser) {
return _.nvgpu.WarpgroupMatrixDescriptorType.parse(parser);
}
parseMBarrierGroup(parser) {
return _.nvgpu.MBarrierGroupType.parse(parser);
}
};
_.NVVMDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'nvvm');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// Reference: NVVMDialect.cpp - parseMmaOperand
// Helper to parse operand list in the format: name[operands]
const parseMmaOperand = (name) => {
parser.parseKeyword(name);
return parser.parseOperandList('optionalSquare');
};
// Reference: NVVMDialect.cpp - MmaOp::parse
if (op === 'nvvm.mma.sync') {
const fragsA = parseMmaOperand('A');
const fragsB = parseMmaOperand('B');
const fragsC = parseMmaOperand('C');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
parser.resolveOperands(fragsA, funcType.inputs, result.operands);
parser.resolveOperands(fragsB, funcType.inputs, result.operands);
parser.resolveOperands(fragsC, funcType.inputs, result.operands);
result.addTypes(funcType.results);
}
}
return true;
}
// Reference: NVVMDialect.cpp - MmaSpOp::parse
if (op === 'nvvm.mma.sp.sync') {
const fragsA = parseMmaOperand('A');
const fragsB = parseMmaOperand('B');
const fragsC = parseMmaOperand('C');
const fragsSparseMetadata = parseMmaOperand('sparseMetadata');
const fragsSelector = parseMmaOperand('selector');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
parser.resolveOperands(fragsA, funcType.inputs, result.operands);
parser.resolveOperands(fragsB, funcType.inputs, result.operands);
parser.resolveOperands(fragsC, funcType.inputs, result.operands);
parser.resolveOperands(fragsSparseMetadata, funcType.inputs, result.operands);
parser.resolveOperands(fragsSelector, funcType.inputs, result.operands);
result.addTypes(funcType.results);
}
}
return true;
}
// Reference: NVVMDialect.cpp - MmaBlockScaleOp::parse
if (op === 'nvvm.mma.block_scale') {
const fragsA = parseMmaOperand('A');
const fragsB = parseMmaOperand('B');
const fragsC = parseMmaOperand('C');
const scaleAOperands = parseMmaOperand('scaleA');
const scaleBOperands = parseMmaOperand('scaleB');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
parser.resolveOperands(fragsA, funcType.inputs, result.operands);
parser.resolveOperands(fragsB, funcType.inputs, result.operands);
parser.resolveOperands(fragsC, funcType.inputs, result.operands);
parser.resolveOperands(scaleAOperands, funcType.inputs, result.operands);
parser.resolveOperands(scaleBOperands, funcType.inputs, result.operands);
result.addTypes(funcType.results);
}
}
return true;
}
// Reference: NVVMDialect.cpp - MmaSpBlockScaleOp::parse
if (op === 'nvvm.mma.sp.block_scale') {
const fragsA = parseMmaOperand('A');
const fragsB = parseMmaOperand('B');
const fragsC = parseMmaOperand('C');
const metadataOperands = parseMmaOperand('sparseMetadata');
const selectorOperands = parseMmaOperand('selector');
const scaleAOperands = parseMmaOperand('scaleA');
const scaleBOperands = parseMmaOperand('scaleB');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
parser.resolveOperands(fragsA, funcType.inputs, result.operands);
parser.resolveOperands(fragsB, funcType.inputs, result.operands);
parser.resolveOperands(fragsC, funcType.inputs, result.operands);
parser.resolveOperands(metadataOperands, funcType.inputs, result.operands);
parser.resolveOperands(selectorOperands, funcType.inputs, result.operands);
parser.resolveOperands(scaleAOperands, funcType.inputs, result.operands);
parser.resolveOperands(scaleBOperands, funcType.inputs, result.operands);
result.addTypes(funcType.results);
}
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.omp = {};
_.omp.MapBoundsType = class extends _.Type {
constructor() {
super(null);
}
static parse(/* parser */) {
return new _.omp.MapBoundsType();
}
toString() {
return '!omp.map_bounds_ty';
}
};
_.omp.AffinityEntryType = class extends _.Type {
constructor(addrType, lenType) {
super(null);
this.addrType = addrType;
this.lenType = lenType;
}
static parse(parser) {
parser.parseLess();
const addrType = parser.parseType();
parser.parseComma();
const lenType = parser.parseType();
parser.parseGreater();
return new _.omp.AffinityEntryType(addrType, lenType);
}
toString() {
return `!omp.affinity_entry_ty<${this.addrType}, ${this.lenType}>`;
}
};
_.omp.IteratedType = class extends _.Type {
constructor(elementType) {
super(null);
this.elementType = elementType;
}
static parse(parser) {
parser.parseLess();
const elementType = parser.parseType();
parser.parseGreater();
return new _.omp.IteratedType(elementType);
}
toString() {
return `!omp.iterated<${this.elementType}>`;
}
};
_.omp.CanonicalLoopInfoType = class extends _.Type {
constructor() {
super(null);
}
static parse(/* parser */) {
return new _.omp.CanonicalLoopInfoType();
}
toString() {
return '!omp.cli';
}
};
_.OpenMPDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'omp');
this.registerCustomDirective('IteratorHeader', this.parseIteratorHeader.bind(this));
this.registerCustomDirective('MapClause', this.parseMapClause.bind(this));
this.registerCustomDirective('CaptureType', this.parseCaptureType.bind(this));
this.registerCustomDirective('MembersIndex', this.parseMembersIndex.bind(this));
this.registerCustomDirective('PrivateReductionRegion', this.parsePrivateReductionRegion.bind(this));
this.registerCustomDirective('PrivateRegion', this.parsePrivateRegion.bind(this));
this.registerCustomDirective('InReductionPrivateRegion', this.parseInReductionPrivateRegion.bind(this));
this.registerCustomDirective('InReductionPrivateReductionRegion', this.parseInReductionPrivateReductionRegion.bind(this));
this.registerCustomDirective('TaskReductionRegion', this.parseTaskReductionRegion.bind(this));
this.registerCustomDirective('UseDeviceAddrUseDevicePtrRegion', this.parseUseDeviceAddrUseDevicePtrRegion.bind(this));
this.registerCustomDirective('TargetOpRegion', this.parseTargetOpRegion.bind(this));
this.registerCustomDirective('ClauseAttr', this.parseClauseAttr.bind(this));
this.registerCustomDirective('DependVarList', this.parseDependVarList.bind(this));
this.registerCustomDirective('LoopTransformClis', this.parseLoopTransformClis.bind(this));
this.registerCustomDirective('SynchronizationHint', this.parseSynchronizationHint.bind(this));
this.registerCustomDirective('AlignedClause', this.parseAlignedClause.bind(this));
this.registerCustomDirective('ScheduleClause', this.parseScheduleClause.bind(this));
this.registerCustomDirective('AllocateAndAllocator', this.parseAllocateAndAllocator.bind(this));
this.registerCustomDirective('DynGroupprivateClause', this.parseDynGroupprivateClause.bind(this));
this.registerCustomDirective('LinearClause', this.parseLinearClause.bind(this));
this.registerCustomDirective('UniformClause', this.parseUniformClause.bind(this));
this.registerCustomDirective('OrderClause', this.parseOrderClause.bind(this));
this.registerCustomDirective('Copyprivate', this.parseCopyprivate.bind(this));
this.registerCustomDirective('GrainsizeClause', this.parseGranularityClause.bind(this));
this.registerCustomDirective('NumTasksClause', this.parseGranularityClause.bind(this));
this.registerCustomDirective('AffinityClause', this.parseAffinityClause.bind(this));
this.registerCustomDirective('HeapAllocClause', this.parseHeapAllocClause.bind(this));
this.registerCustomAttribute('DataSharingClauseTypeAttr', this.parseDataSharingClauseTypeAttr.bind(this));
this.registerCustomAttribute('ClauseCancelConstructTypeAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('ClauseDependAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('ClauseOrderingIncludeTypeAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('ClauseTypeAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('ClauseDistScheduleTypeAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('OrderModifierAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('DeclareTargetDeviceTypeAttr', this.parseParenthesizedEnumAttr.bind(this));
this.registerCustomAttribute('TargetExecModeAttr', this.parseParenthesizedEnumAttr.bind(this));
}
parseType(parser) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'map_bounds_ty') {
return _.omp.MapBoundsType.parse(parser);
}
if (mnemonic === 'affinity_entry_ty') {
return _.omp.AffinityEntryType.parse(parser);
}
if (mnemonic === 'iterated') {
return _.omp.IteratedType.parse(parser);
}
if (mnemonic === 'cli') {
return _.omp.CanonicalLoopInfoType.parse(parser);
}
}
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'omp.loop_nest') {
return this.parseLoopNestOp(parser, result);
}
if (op === 'omp.canonical_loop') {
return this.parseCanonicalLoopOp(parser, result);
}
if (op === 'omp.unroll_heuristic') {
return this.parseUnrollHeuristicOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseCanonicalLoopOp(parser, result) {
if (parser.parseOptionalLParen()) {
const cliOperand = parser.parseOperand();
// CLI operand is a loop handle, resolve with null type
parser.resolveOperand(cliOperand, null, result.operands);
parser.parseRParen();
}
const inductionVar = parser.parseOperand();
parser.parseColon();
const ivType = parser.parseType();
parser.parseKeyword('in');
parser.parseKeyword('range');
parser.parseLParen();
const rangeOperand = parser.parseOperand();
parser.resolveOperand(rangeOperand, null, result.operands);
parser.parseRParen();
const region = result.addRegion();
// Pass induction variable as region argument
const regionArgs = [{ name: inductionVar.name, type: ivType }];
parser.parseRegion(region, regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseUnrollHeuristicOp(parser, result) {
parser.parseLParen();
const applyee = parser.parseOperand();
// Applyee is a loop handle, resolve with null type
parser.resolveOperand(applyee, null, result.operands);
parser.parseRParen();
if (parser.parseOptionalArrow()) {
parser.parseLParen();
parser.parseRParen();
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseLoopNestOp(parser, result) {
// Parse CLI operands (loop handles)
const unresolvedCli = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
unresolvedCli.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
// Parse types for CLI operands — C++ uses parseColonType (single type)
const cliTypes = [];
if (parser.parseOptionalColon()) {
cliTypes.push(parser.parseType());
}
parser.resolveOperands(unresolvedCli, cliTypes, result.operands);
if (parser.parseOptionalEqual()) {
if (parser.parseOptionalLParen()) {
const unresolvedLb = [];
if (!parser.parseOptionalRParen()) {
do {
unresolvedLb.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
for (const lb of unresolvedLb) {
parser.resolveOperand(lb, null, result.operands);
}
}
if (parser.parseOptionalKeyword('to')) {
if (parser.parseOptionalLParen()) {
const unresolvedUb = [];
if (!parser.parseOptionalRParen()) {
do {
unresolvedUb.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
for (const ub of unresolvedUb) {
parser.resolveOperand(ub, null, result.operands);
}
}
}
parser.parseOptionalKeyword('inclusive');
if (parser.parseOptionalKeyword('step')) {
if (parser.parseOptionalLParen()) {
const unresolvedStep = [];
if (!parser.parseOptionalRParen()) {
do {
unresolvedStep.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
for (const step of unresolvedStep) {
parser.resolveOperand(step, null, result.operands);
}
}
}
}
if (parser.parseOptionalKeyword('collapse')) {
parser.parseLParen();
const value = parser.parseInteger();
result.addAttribute('collapse_num_loops', value);
parser.parseRParen();
}
if (parser.parseOptionalKeyword('tiles')) {
parser.parseLParen();
const tiles = [];
if (!parser.parseOptionalRParen()) {
do {
tiles.push(parser.parseInteger());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
result.addAttribute('tile_sizes', tiles);
}
const region = result.addRegion();
parser.parseRegion(region);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseParenthesizedEnumAttr(parser) {
if (parser.parseOptionalLParen()) {
const value = parser.parseOptionalKeyword();
parser.parseRParen();
return new _.TypedAttr(value, null);
}
return null;
}
parseOrderClause(parser, op, orderAttrName, orderModAttrName) {
const orderModifiers = ['reproducible', 'unconstrained'];
let keyword = parser.parseKeyword();
if (orderModifiers.includes(keyword)) {
op.addAttribute(orderModAttrName, keyword);
parser.parseColon();
keyword = parser.parseKeyword();
}
op.addAttribute(orderAttrName, keyword);
}
parseLinearClause(parser, op, linearVars, linearTypes, linearStepVars, linearStepTypes, linearModifiersAttr) {
const modifiers = [];
parser.parseCommaSeparatedList('none', () => {
const modifier = parser.parseOptionalKeyword(['val', 'ref', 'uval']);
if (modifier) {
parser.parseLParen();
}
linearVars.push(parser.parseOperand());
parser.parseColon();
linearTypes.push(parser.parseType());
parser.parseEqual();
linearStepVars.push(parser.parseOperand());
parser.parseColon();
linearStepTypes.push(parser.parseType());
if (modifier) {
parser.parseRParen();
}
modifiers.push(modifier || null);
});
if (modifiers.some((m) => m !== null)) {
op.addAttribute(linearModifiersAttr, modifiers);
}
parser.resolveOperands(linearVars, linearTypes, op.operands);
parser.resolveOperands(linearStepVars, linearStepTypes, op.operands);
}
parseUniformClause(parser, result, uniformVars, uniformTypes) {
parser.parseCommaSeparatedList('none', () => {
uniformVars.push(parser.parseOperand());
parser.parseColon();
uniformTypes.push(parser.parseType());
});
}
parseCopyprivate(parser, op, varsAttr, typesAttr, symsAttr) {
const unresolvedVars = [];
const varTypes = [];
const copyprivateSyms = [];
do {
unresolvedVars.push(parser.parseOperand());
parser.parseArrow();
const sym = parser.parseOptionalSymbolName();
parser.parseColon();
const type = parser.parseType();
varTypes.push(type);
copyprivateSyms.push(sym);
} while (parser.parseOptionalComma());
parser.resolveOperands(unresolvedVars, varTypes, op.operands);
if (symsAttr) {
op.addAttribute(symsAttr, copyprivateSyms);
}
}
parseGranularityClause(parser, op, modAttr) {
const modifier = parser.parseOptionalKeyword();
if (modifier) {
parser.parseComma();
}
const unresolvedOperand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(unresolvedOperand, type, op.operands);
if (modAttr && modifier) {
op.addAttribute(modAttr, modifier);
}
}
parseAlignedClause(parser, op, alignedVars, alignedTypes, alignmentsAttrName) {
const alignments = [];
parser.parseCommaSeparatedList('none', () => {
alignedVars.push(parser.parseOperand());
parser.parseColon();
alignedTypes.push(parser.parseType());
parser.parseArrow();
alignments.push(parser.parseAttribute());
});
if (alignments.length > 0) {
op.addAttribute(alignmentsAttrName, alignments);
}
}
// Following reference: parseAffinityClause / parseSplitIteratedList in OpenMPDialect.cpp:4658/1438
parseAffinityClause(parser, result, iterated, affinityVars, iteratedTypes, affinityTypes) {
parser.parseCommaSeparatedList('none', () => {
const operand = parser.parseOperand();
const type = parser.parseColonType();
if (type instanceof _.omp.IteratedType) {
iterated.push(operand);
iteratedTypes.push(type);
} else {
affinityVars.push(operand);
affinityTypes.push(type);
}
});
}
parseHeapAllocClause(parser, op, inTypeAttrName, typeparamsOps, typeparamsTypes, shapeOps, shapeTypes) {
const inType = parser.parseType();
op.addAttribute(inTypeAttrName, { value: inType, type: 'type' });
if (parser.parseOptionalLParen()) {
do {
typeparamsOps.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseColon();
for (const t of parser.parseTypeList()) {
typeparamsTypes.push(t);
}
parser.parseRParen();
}
if (parser.parseOptionalComma()) {
do {
shapeOps.push(parser.parseOperand());
} while (parser.parseOptionalComma());
const indexType = new _.IndexType();
for (let i = 0; i < shapeOps.length; i++) {
shapeTypes.push(indexType);
}
}
}
parseScheduleClause(parser, op, scheduleKindAttrName, scheduleModAttrName, scheduleSimdAttrName, chunkOperands, chunkTypes) {
const keyword = parser.parseKeyword();
op.addAttribute(scheduleKindAttrName, keyword);
if (keyword === 'static' || keyword === 'dynamic' || keyword === 'guided') {
if (parser.parseOptionalEqual()) {
chunkOperands.push(parser.parseOperand());
parser.parseColon();
chunkTypes.push(parser.parseType());
}
}
const modifiers = [];
while (parser.parseOptionalComma()) {
modifiers.push(parser.parseKeyword());
}
if (modifiers.length > 0) {
op.addAttribute(scheduleModAttrName, modifiers[0]);
if (modifiers.length > 1) {
op.addAttribute(scheduleSimdAttrName, new _.UnitAttr());
}
}
}
parseAllocateAndAllocator(parser, op, allocateVars, allocateTypes, allocatorVars, allocatorTypes) {
parser.parseCommaSeparatedList('none', () => {
allocatorVars.push(parser.parseOperand());
parser.parseColon();
allocatorTypes.push(parser.parseType());
parser.parseArrow();
allocateVars.push(parser.parseOperand());
parser.parseColon();
allocateTypes.push(parser.parseType());
});
}
parseDynGroupprivateClause(parser, op, accessGroupAttrName, fallbackAttrName, sizeOperands, sizeTypes) {
parser.parseCommaSeparatedList('none', () => {
if (parser.parseOptionalKeyword('cgroup')) {
op.addAttribute(accessGroupAttrName, 'cgroup');
return;
}
if (parser.parseOptionalKeyword('fallback')) {
parser.parseLParen();
const fbKind = parser.parseKeyword();
op.addAttribute(fallbackAttrName, fbKind);
parser.parseRParen();
return;
}
sizeOperands.push(parser.parseOperand());
parser.parseColon();
sizeTypes.push(parser.parseType());
});
}
parseSynchronizationHint(parser, op, hintAttr = 'hint') {
if (parser.parseOptionalKeyword('none')) {
op.addAttribute(hintAttr, 0);
return;
}
let hint = 0;
const hints = [];
let keyword = parser.parseOptionalKeyword();
while (keyword) {
hints.push(keyword);
if (keyword === 'uncontended') {
hint |= 1;
} else if (keyword === 'contended') {
hint |= 2;
} else if (keyword === 'nonspeculative') {
hint |= 4;
} else if (keyword === 'speculative') {
hint |= 8;
}
if (!parser.parseOptionalComma()) {
break;
}
keyword = parser.parseOptionalKeyword();
}
op.addAttribute(hintAttr, hint);
}
parseClauseAttr(parser, op, attrName) {
const enumValue = parser.parseOptionalKeyword();
if (enumValue && attrName) {
op.addAttribute(attrName, enumValue);
}
}
// Parses: keyword([byref] [@sym] %operand -> %block_arg [map_idx = N] : type, ...) [private_barrier]
// Mirrors parseClauseWithRegionArgs / parseBlockArgClause from OpenMPDialect.cpp
parseBlockArgClause(parser, entryBlockArgs, keyword, vars, types, clauseConfig) {
if (!parser.parseOptionalKeyword(keyword)) {
return;
}
if (!vars) {
return;
}
parser.parseLParen();
const syms = [];
const byrefFlags = [];
const mapIndices = [];
// Optional reduction modifier: mod : <keyword> ,
if (clauseConfig.modifier && parser.parseOptionalKeyword('mod')) {
parser.parseColon();
const modValue = parser.parseKeyword();
parser.parseComma();
clauseConfig.modifier.value = modValue;
}
if (!parser.parseOptionalRParen()) {
do {
if (clauseConfig.byref) {
byrefFlags.push(Boolean(parser.parseOptionalKeyword('byref')));
}
if (clauseConfig.syms) {
syms.push(parser.parseOptionalSymbolName());
}
const operand = parser.parseOperand();
vars.push(operand);
parser.parseArrow();
const blockArg = parser.parseOperand();
entryBlockArgs.push(blockArg);
if (clauseConfig.mapIndices) {
if (parser.parseOptionalLSquare()) {
parser.parseKeyword('map_idx');
parser.parseEqual();
mapIndices.push(parser.parseInteger());
parser.parseRSquare();
} else {
mapIndices.push(-1);
}
}
} while (parser.parseOptionalComma());
parser.parseColon();
do {
types.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
if (clauseConfig.needsBarrier && parser.parseOptionalKeyword('private_barrier')) {
clauseConfig.needsBarrier.value = true;
}
if (clauseConfig.syms && syms.length > 0) {
clauseConfig.syms.value = syms;
}
if (clauseConfig.byref && byrefFlags.length > 0) {
clauseConfig.byref.value = byrefFlags;
}
if (clauseConfig.mapIndices && mapIndices.length > 0 && mapIndices.some((v) => v !== -1)) {
clauseConfig.mapIndices.value = mapIndices;
}
}
// Common helper: parses block arg clauses then the region.
// clauseArgs is an array of { keyword, vars, types, config } objects.
parseBlockArgRegion(parser, op, regionName, clauseArgs) {
const entryBlockArgs = [];
for (const clause of clauseArgs) {
this.parseBlockArgClause(parser, entryBlockArgs, clause.keyword, clause.vars, clause.types, clause.config);
}
// Set entry block arg types
let typeIdx = 0;
for (const clause of clauseArgs) {
if (clause.types && clause.types.length > 0) {
for (let i = 0; i < clause.types.length && typeIdx < entryBlockArgs.length; i++, typeIdx++) {
entryBlockArgs[typeIdx].type = clause.types[i];
}
}
}
const regionArgs = entryBlockArgs.map((arg) => ({
name: arg.name,
type: arg.type || null
}));
const region = op.addRegion();
parser.parseRegion(region, regionArgs.length > 0 ? regionArgs : undefined);
for (const clause of clauseArgs) {
const cfg = clause.config;
if (cfg.syms && cfg.syms.value && cfg.syms.attrName) {
op.addAttribute(cfg.syms.attrName, cfg.syms.value);
}
if (cfg.byref && cfg.byref.value && cfg.byref.attrName) {
op.addAttribute(cfg.byref.attrName, cfg.byref.value);
}
if (cfg.needsBarrier && cfg.needsBarrier.value && cfg.needsBarrier.attrName) {
op.addAttribute(cfg.needsBarrier.attrName, true);
}
if (cfg.modifier && cfg.modifier.value && cfg.modifier.attrName) {
op.addAttribute(cfg.modifier.attrName, cfg.modifier.value);
}
if (cfg.mapIndices && cfg.mapIndices.value && cfg.mapIndices.attrName) {
op.addAttribute(cfg.mapIndices.attrName, cfg.mapIndices.value);
}
}
}
parseTargetOpRegion(parser, op, regionName, hasDeviceAddrVars, hasDeviceAddrTypes,hostEvalVars, hostEvalTypes, inReductionVars, inReductionTypes, inReductionByrefAttr, inReductionSymsAttr, mapVars, mapTypes, privateVars, privateTypes, privateSymsAttr, privateNeedsBarrierAttr, privateMapsAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'has_device_addr', vars: hasDeviceAddrVars, types: hasDeviceAddrTypes, config: {} },
{ keyword: 'host_eval', vars: hostEvalVars, types: hostEvalTypes, config: {} },
{ keyword: 'in_reduction', vars: inReductionVars, types: inReductionTypes, config: { syms: { attrName: inReductionSymsAttr }, byref: { attrName: inReductionByrefAttr } } },
{ keyword: 'map_entries', vars: mapVars, types: mapTypes, config: {} },
{ keyword: 'private', vars: privateVars, types: privateTypes, config: { syms: { attrName: privateSymsAttr }, needsBarrier: { attrName: privateNeedsBarrierAttr }, mapIndices: { attrName: privateMapsAttr } } },
]);
}
parseIteratorHeader(parser, op) {
const regionArgs = [];
const lowerBounds = [];
const upperBounds = [];
const steps = [];
do {
const arg = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
regionArgs.push({ name: arg.name, type });
} while (parser.parseOptionalComma());
parser.parseRParen();
parser.parseEqual();
parser.parseLParen();
do {
lowerBounds.push(parser.parseOperand());
parser.parseKeyword('to');
upperBounds.push(parser.parseOperand());
parser.parseKeyword('step');
steps.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
const region = op.addRegion();
parser.parseRegion(region, regionArgs);
const indexType = new _.IndexType();
for (const lb of lowerBounds) {
parser.resolveOperand(lb, indexType, op.operands);
}
for (const ub of upperBounds) {
parser.resolveOperand(ub, indexType, op.operands);
}
for (const st of steps) {
parser.resolveOperand(st, indexType, op.operands);
}
}
parseMapClause(parser, op, attrName = 'map_type') {
const mapFlags = [];
do {
const flag = parser.parseOptionalKeyword();
if (flag) {
mapFlags.push(flag);
}
} while (parser.parseOptionalComma());
if (attrName && mapFlags.length > 0) {
op.addAttribute(attrName, mapFlags.join(', '));
}
}
parseCaptureType(parser, op, attrName) {
const captureType = parser.parseOptionalKeyword();
if (captureType) {
if (attrName) {
op.addAttribute(attrName, captureType);
}
}
}
parseMembersIndex(parser, op, attrName) {
const memberIndices = [];
do {
if (parser.parseOptionalLSquare()) {
const indices = [];
do {
const idx = parser.parseOptionalInteger();
if (idx !== null) {
indices.push(idx);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
memberIndices.push(indices);
}
} while (parser.parseOptionalComma());
if (attrName && memberIndices.length > 0) {
op.addAttribute(attrName, memberIndices);
}
}
parsePrivateReductionRegion(parser, op, regionName, privateVars, privateTypes, privateSymsAttr, privateNeedsBarrierAttr, reductionModAttr, reductionVars, reductionTypes, reductionByrefAttr, reductionSymsAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'private', vars: privateVars, types: privateTypes, config: { syms: { attrName: privateSymsAttr }, needsBarrier: { attrName: privateNeedsBarrierAttr } } },
{ keyword: 'reduction', vars: reductionVars, types: reductionTypes, config: { syms: { attrName: reductionSymsAttr }, byref: { attrName: reductionByrefAttr }, modifier: { attrName: reductionModAttr } } },
]);
}
parsePrivateRegion(parser, op, regionName, privateVars, privateTypes, privateSymsAttr, privateNeedsBarrierAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'private', vars: privateVars, types: privateTypes, config: { syms: { attrName: privateSymsAttr }, needsBarrier: { attrName: privateNeedsBarrierAttr } } },
]);
}
parseInReductionPrivateRegion(parser, op, regionName, inReductionVars, inReductionTypes, inReductionByrefAttr, inReductionSymsAttr, privateVars, privateTypes, privateSymsAttr, privateNeedsBarrierAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'in_reduction', vars: inReductionVars, types: inReductionTypes, config: { syms: { attrName: inReductionSymsAttr }, byref: { attrName: inReductionByrefAttr } } },
{ keyword: 'private', vars: privateVars, types: privateTypes, config: { syms: { attrName: privateSymsAttr }, needsBarrier: { attrName: privateNeedsBarrierAttr } } },
]);
}
parseInReductionPrivateReductionRegion(parser, op, regionName, inReductionVars, inReductionTypes, inReductionByrefAttr, inReductionSymsAttr, privateVars, privateTypes, privateSymsAttr, privateNeedsBarrierAttr, reductionModAttr, reductionVars, reductionTypes, reductionByrefAttr, reductionSymsAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'in_reduction', vars: inReductionVars, types: inReductionTypes, config: { syms: { attrName: inReductionSymsAttr }, byref: { attrName: inReductionByrefAttr } } },
{ keyword: 'private', vars: privateVars, types: privateTypes, config: { syms: { attrName: privateSymsAttr }, needsBarrier: { attrName: privateNeedsBarrierAttr } } },
{ keyword: 'reduction', vars: reductionVars, types: reductionTypes, config: { syms: { attrName: reductionSymsAttr }, byref: { attrName: reductionByrefAttr }, modifier: { attrName: reductionModAttr } } },
]);
}
parseTaskReductionRegion(parser, op, regionName, taskReductionVars, taskReductionTypes, taskReductionByrefAttr, taskReductionSymsAttr) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'task_reduction', vars: taskReductionVars, types: taskReductionTypes, config: { syms: { attrName: taskReductionSymsAttr }, byref: { attrName: taskReductionByrefAttr } } },
]);
}
parseUseDeviceAddrUseDevicePtrRegion(parser, op, regionName, useDeviceAddrVars, useDeviceAddrTypes, useDevicePtrVars, useDevicePtrTypes) {
this.parseBlockArgRegion(parser, op, regionName, [
{ keyword: 'use_device_addr', vars: useDeviceAddrVars, types: useDeviceAddrTypes, config: {} },
{ keyword: 'use_device_ptr', vars: useDevicePtrVars, types: useDevicePtrTypes, config: {} },
]);
}
parseDataSharingClauseTypeAttr(parser) {
if (parser.parseOptionalLBrace()) {
parser.parseKeyword('type');
parser.parseEqual();
const value = parser.parseKeyword();
parser.parseRBrace();
return { value };
}
return null;
}
parseDependVarList(parser, op, operandAttr, typesAttr, kindAttr, iteratedOperandAttr, iteratedTypesAttr, iteratedKindAttr) {
const dependVars = [];
const dependTypes = [];
const dependKinds = [];
const iteratedVars = [];
const iteratedTypes = [];
const iteratedKinds = [];
do {
const keyword = parser.parseKeyword();
parser.parseArrow();
const operand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
if (iteratedOperandAttr && type instanceof _.omp.IteratedType) {
iteratedKinds.push(keyword);
iteratedVars.push(operand);
iteratedTypes.push(type);
} else {
dependKinds.push(keyword);
dependVars.push(operand);
dependTypes.push(type);
}
} while (parser.parseOptionalComma());
if (operandAttr) {
for (let i = 0; i < dependVars.length; i++) {
const type = i < dependTypes.length ? dependTypes[i] : null;
parser.resolveOperand(dependVars[i], type, op.operands);
}
}
if (kindAttr) {
op.addAttribute(kindAttr, dependKinds);
}
if (iteratedOperandAttr) {
for (let i = 0; i < iteratedVars.length; i++) {
const type = i < iteratedTypes.length ? iteratedTypes[i] : null;
parser.resolveOperand(iteratedVars[i], type, op.operands);
}
}
if (iteratedKindAttr) {
op.addAttribute(iteratedKindAttr, iteratedKinds);
}
}
parseLoopTransformClis(parser, op, generateesOperands, applyeesOperands) {
if (!parser.parseOptionalLess()) {
generateesOperands.push(...parser.parseOperandList('paren'));
parser.parseLess();
}
parser.parseMinus();
applyeesOperands.push(...parser.parseOperandList('paren'));
}
};
_.llvm = {};
_.llvm.LLVMPointerType = class extends _.Type {
constructor(addressSpace, pointeeType) {
super(null);
this.addressSpace = addressSpace || 0;
this.pointeeType = pointeeType || null;
}
static parse(parser) {
if (parser.parseOptionalLess()) {
const addressSpaceOrType = parser.parseOptionalInteger();
if (addressSpaceOrType !== null) {
parser.parseGreater();
return new _.llvm.LLVMPointerType(addressSpaceOrType);
}
const pointee = _.llvm.LLVMDialect.dispatchParse(parser);
parser.parseGreater();
return new _.llvm.LLVMPointerType(0, pointee);
}
return new _.llvm.LLVMPointerType(0);
}
toString() {
if (this.pointeeType) {
return `!llvm.ptr<${this.pointeeType}>`;
}
if (this.addressSpace) {
return `!llvm.ptr<${this.addressSpace}>`;
}
return '!llvm.ptr';
}
};
_.llvm.LLVMArrayType = class extends _.Type {
constructor(elementType, numElements) {
super(null);
this.elementType = elementType;
this.numElements = numElements;
}
static parse(parser) {
parser.parseLess();
const numElements = parser.parseInteger();
parser.parseKeyword('x');
const elementType = _.llvm.LLVMDialect.dispatchParse(parser);
parser.parseGreater();
return new _.llvm.LLVMArrayType(elementType, numElements);
}
toString() {
return `!llvm.array<${this.numElements} x ${this.elementType}>`;
}
};
_.llvm.LLVMStructType = class extends _.Type {
constructor(name, body, isPacked) {
super(null);
this.structName = name;
this.body = body;
this.isPacked = isPacked || false;
}
static parse(parser) {
parser.parseLess();
let name = null;
if (parser.parseOptionalString) {
name = parser.parseOptionalString();
}
if (name !== null) {
if (parser.parseOptionalGreater()) {
return new _.llvm.LLVMStructType(name, null, false);
}
parser.parseComma();
}
if (parser.parseOptionalKeyword('opaque')) {
parser.parseGreater();
return new _.llvm.LLVMStructType(name, null, false);
}
const isPacked = Boolean(parser.parseOptionalKeyword('packed'));
parser.parseLParen();
const body = [];
if (!parser.parseOptionalRParen()) {
body.push(_.llvm.LLVMDialect.dispatchParse(parser));
while (parser.parseOptionalComma()) {
body.push(_.llvm.LLVMDialect.dispatchParse(parser));
}
parser.parseRParen();
}
parser.parseGreater();
return new _.llvm.LLVMStructType(name, body, isPacked);
}
toString() {
const parts = [];
if (this.structName) {
parts.push(`"${this.structName}"`);
}
if (this.body === null) {
if (this.structName) {
return `!llvm.struct<"${this.structName}">`;
}
return '!llvm.struct<opaque>';
}
const prefix = this.isPacked ? 'packed ' : '';
const members = this.body.map((t) => t?.toString ? t.toString() : t).join(', ');
if (this.structName) {
return `!llvm.struct<"${this.structName}", ${prefix}(${members})>`;
}
return `!llvm.struct<${prefix}(${members})>`;
}
};
_.llvm.LLVMFunctionType = class extends _.Type {
constructor(returnType, params, varArg = false) {
super(null);
this.returnType = returnType;
this.params = params || [];
this.varArg = varArg;
}
static parse(parser) {
parser.parseLess();
const returnType = _.llvm.LLVMDialect.dispatchParse(parser);
parser.parseLParen();
const params = [];
let varArg = false;
if (!parser.parseOptionalRParen()) {
if (parser.parseOptionalEllipsis()) {
varArg = true;
} else {
params.push(_.llvm.LLVMDialect.dispatchParse(parser));
while (parser.parseOptionalComma()) {
if (parser.parseOptionalEllipsis()) {
varArg = true;
break;
}
params.push(_.llvm.LLVMDialect.dispatchParse(parser));
}
}
parser.parseRParen();
}
parser.parseGreater();
return new _.llvm.LLVMFunctionType(returnType, params, varArg);
}
get inputs() {
return this.params;
}
get results() {
return this.returnType ? [this.returnType] : [];
}
toString() {
const params = this.params.map((t) => t.toString());
if (this.varArg) {
params.push('...');
}
const returnType = this.returnType ? this.returnType.toString() : 'void';
return `!llvm.func<${returnType} (${params.join(', ')})>`;
}
};
_.llvm.LLVMTargetExtType = class extends _.Type {
constructor(extTypeName, typeParams, intParams) {
super(null);
this.extTypeName = extTypeName;
this.typeParams = typeParams;
this.intParams = intParams;
}
static parse(parser) {
parser.parseLess();
const extTypeName = parser.parseString();
const typeParams = [];
const intParams = [];
let parseType = true;
while (parser.parseOptionalComma()) {
const intVal = parser.parseOptionalInteger();
if (intVal !== null) {
intParams.push(intVal);
parseType = false;
} else if (parseType) {
typeParams.push(parser.parseType());
} else {
parser.emitError('Expected integer parameter');
}
}
parser.parseGreater();
return new _.llvm.LLVMTargetExtType(extTypeName, typeParams, intParams);
}
toString() {
const parts = [this.extTypeName];
for (const t of this.typeParams) {
parts.push(t.toString());
}
for (const i of this.intParams) {
parts.push(String(i));
}
return `!llvm.target<${parts.join(', ')}>`;
}
};
_.llvm.LLVMDialect = class extends _.Dialect {
constructor(operations, name = 'llvm') {
super(operations, name);
this.registerCustomDirective('GEPIndices', this.parseGEPIndices.bind(this));
this.registerCustomDirective('IndirectBrOpSucessors', this.parseIndirectBrOpSucessors.bind(this));
this.registerCustomDirective('InsertExtractValueElementType', this.parseInsertExtractValueElementType.bind(this));
this.registerCustomDirective('LLVMLinkage', this.parseLLVMLinkage.bind(this));
this.registerCustomDirective('OpBundles', this.parseOpBundles.bind(this));
this.registerCustomDirective('ShuffleType', this.parseShuffleType.bind(this));
this.registerCustomDirective('SwitchOpCases', this.parseSwitchOpCases.bind(this));
this.registerCustomAttribute('LLVM_IntegerOverflowFlagsProp', this.parseLLVMIntegerOverflowFlagsProp.bind(this));
this.registerCustomAttribute('GEPNoWrapFlagsProp', (parser, type) => this.parseEnumFlags(parser, type, '|', true));
this.registerCustomAttribute('LLVM_BlockAddressAttr', this.parseLLVMBlockAddressAttr.bind(this));
this.registerCustomAttribute('LLVM_BlockTagAttr', this.parseLLVMBlockTagAttr.bind(this));
this.registerCustomType('LLVM_AnyPointer', this.parseLLVMPointerType.bind(this));
this.registerCustomType('LLVM_PointerInAddressSpace', this.parseLLVMPointerType.bind(this));
this.registerCustomType('LLVM_Type', (parser) => parser.parseType());
}
static dispatchParse(parser) {
const type = parser.parseOptionalType();
if (type) {
return type;
}
const key = parser.parseKeyword();
switch (key) {
case 'void': return null;
case 'ppc_fp128': return new _.Type('!llvm.ppc_fp128');
case 'token': return new _.Type('!llvm.token');
case 'label': return new _.Type('!llvm.label');
case 'metadata': return new _.Type('!llvm.metadata');
case 'func': return _.llvm.LLVMFunctionType.parse(parser);
case 'ptr': return _.llvm.LLVMPointerType.parse(parser);
case 'array': return _.llvm.LLVMArrayType.parse(parser);
case 'struct': return _.llvm.LLVMStructType.parse(parser);
case 'target': return _.llvm.LLVMTargetExtType.parse(parser);
case 'x86_amx': return new _.Type('!llvm.x86_amx');
default: throw new mlir.Error(`Unknown LLVM type '${key}'.`);
}
}
parseLLVMIntegerOverflowFlagsProp(parser) {
if (parser.parseOptionalKeyword('overflow')) {
return this.parseEnumFlagsAngleBracketComma(parser, { values: ['wrap', 'nuw', 'nsw'] });
}
return null;
}
parseLLVMBlockAddressAttr(parser) {
parser.parseLess();
const result = {};
do {
const key = parser.parseKeyword();
parser.parseEqual();
if (key === 'function') {
result.function = parser.parseAttribute();
} else if (key === 'tag') {
result.tag = this.parseLLVMBlockTagAttr(parser);
}
} while (parser.parseOptionalComma());
parser.parseGreater();
return result;
}
parseLLVMBlockTagAttr(parser) {
parser.parseLess();
const result = {};
do {
const key = parser.parseKeyword();
parser.parseEqual();
if (key === 'id') {
result.id = parser.parseInteger();
}
} while (parser.parseOptionalComma());
parser.parseGreater();
return result;
}
parseLLVMPointerType(parser) {
const type = parser.parseOptionalType();
if (type) {
return type;
}
return _.llvm.LLVMPointerType.parse(parser);
}
getInsertExtractValueElementType(containerType, position) {
let llvmType = containerType.toString ? containerType.toString() : String(containerType);
for (const idx of position) {
const arrayMatch = llvmType.match(/^!llvm\.array<(\d+)\s*x\s*(.+)>$/);
if (arrayMatch) {
llvmType = arrayMatch[2].trim();
continue;
}
const structMatch = llvmType.match(/^!llvm\.struct<(?:"[^"]*",\s*)?\((.+)\)>$/);
if (structMatch) {
const elementsStr = structMatch[1];
const body = [];
let depth = 0;
let current = '';
for (const char of elementsStr) {
if (char === '<' || char === '(' || char === '[' || char === '{') {
depth++;
current += char;
} else if (char === '>' || char === ')' || char === ']' || char === '}') {
depth--;
current += char;
} else if (char === ',' && depth === 0) {
body.push(current.trim());
current = '';
} else {
current += char;
}
}
if (current.trim()) {
body.push(current.trim());
}
if (idx >= 0 && idx < body.length) {
llvmType = body[idx];
continue;
}
}
break;
}
return llvmType;
}
parseInsertExtractValueElementType(parser, op, resultTypes, containerTypes, positionArg) {
// Infer the value type from the container type and position.
// resultTypes = op.types (array to populate)
// containerTypes = array of container types from parsing
// positionArg = position attribute value (ArrayAttr or array)
if (containerTypes && containerTypes.length > 0) {
const containerType = containerTypes[0];
// Use the position argument if provided, otherwise fall back to op.attributes
const positionAttr = positionArg === undefined ? op.attributes.get('position') : positionArg;
const position = positionAttr && positionAttr.value ? positionAttr.value : positionAttr;
if (position && Array.isArray(position) && position.length > 0) {
const valueType = this.getInsertExtractValueElementType(containerType, position);
if (valueType && resultTypes.length === 0) {
resultTypes.push(new _.Type(valueType));
}
}
}
}
parseLLVMLinkage(parser, op /*, args */) {
const linkage = parser.parseOptionalKeyword();
if (linkage) {
op.addAttribute('linkage', linkage);
}
}
parseOpBundles(parser, op /*, args */) {
// Parse operation bundles: [] or ["tag"()] or ["tag"(%0, %1 : i32, i32), ...]
// Returns: null if not present, true if success, throws on failure
// args[0] = $op_bundle_operands - operands for bundles
// args[1] = type($op_bundle_operands) - types
// args[2] = $op_bundle_tags - tags attribute
if (!parser.parseOptionalLSquare()) {
return null; // Not present (equivalent to std::nullopt)
}
if (parser.parseOptionalRSquare()) {
return true; // Success
}
const opBundles = [];
do {
const tag = parser.parseString();
parser.parseLParen();
const bundleOperands = [];
const bundleTypes = [];
if (!parser.parseOptionalRParen()) {
do {
bundleOperands.push(parser.parseAttribute());
} while (parser.parseOptionalComma());
parser.parseColon();
do {
bundleTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
opBundles.push({ tag, operands: bundleOperands, types: bundleTypes });
} while (parser.parseOptionalComma());
parser.parseRSquare();
if (opBundles.length > 0) {
op.addAttribute('op_bundle_tags', opBundles);
}
return true; // Success
}
parseShuffleType(parser, op, v1Types, resTypes, maskAttr) {
if (resTypes.length > 0) {
return;
}
if (!v1Types || v1Types.length === 0 || !maskAttr) {
return;
}
const v1Type = v1Types[0];
if (!(v1Type instanceof _.VectorType)) {
return;
}
const mask = maskAttr.value === undefined ? maskAttr : maskAttr.value;
const maskLen = Array.isArray(mask) ? mask.length : 0;
if (maskLen > 0) {
const isScalable = v1Type.scalableDims && v1Type.scalableDims.some(Boolean);
const resType = new _.VectorType([BigInt(maskLen)], v1Type.elementType, isScalable ? [true] : []);
op.addTypes([resType]);
}
}
// Parse switch operation cases
// Where case: integer `:` bb-id (`(` ssa-use-and-type-list `)`)?
parseSwitchOpCases(parser, op /*, args */) {
// args[0] is ref(type($value)) - the flag type
// args[1] is $case_values - attribute to populate
// args[2] is $caseDestinations - successors array
// args[3] is $caseOperands - operands for each case
// args[4] is type($caseOperands) - types for case operands
if (!parser.parseOptionalLSquare()) {
return;
}
if (parser.parseOptionalRSquare()) {
return;
}
const caseValues = [];
const caseDestinations = [];
const caseOperands = [];
do {
// Handle negative case values: -1, -2, etc.
let sign = 1;
if (parser.parseOptionalMinus()) {
sign = -1;
}
const value = sign * parser.parseInteger();
caseValues.push(value);
parser.parseColon();
const successor = parser.parseOptionalSuccessor();
caseDestinations.push(successor);
if (parser.parseOptionalLParen()) {
const operands = [];
do {
const caseOp = parser.parseOptionalOperand();
if (caseOp) {
operands.push({ name: caseOp });
} else {
break;
}
} while (parser.parseOptionalComma());
if (parser.parseOptionalColon()) {
let idx = 0;
while (idx < operands.length) {
const type = parser.parseType();
if (operands[idx]) {
operands[idx].type = type;
}
idx++;
parser.parseOptionalComma();
}
}
parser.parseRParen();
caseOperands.push(operands);
} else {
caseOperands.push([]);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
if (caseValues.length > 0) {
op.addAttribute('case_values', caseValues);
}
if (caseDestinations.length > 0) {
if (!op.successors) {
op.successors = [];
}
for (const dest of caseDestinations) {
op.successors.push({ name: dest });
}
}
// Note: caseOperands handling would require more complex logic
// to properly associate operands with their successors
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'struct') {
return _.llvm.LLVMStructType.parse(parser);
}
if (mnemonic === 'func') {
return _.llvm.LLVMFunctionType.parse(parser);
}
if (mnemonic === 'ptr') {
return _.llvm.LLVMPointerType.parse(parser);
}
if (mnemonic === 'array') {
return _.llvm.LLVMArrayType.parse(parser);
}
if (mnemonic === 'target') {
return _.llvm.LLVMTargetExtType.parse(parser);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'llvm.func') {
return this.parseLLVMFuncOp(parser, result);
}
if (op === 'llvm.mlir.global') {
return this.parseLLVMGlobalOp(parser, result);
}
if (op === 'llvm.mlir.alias') {
return this.parseLLVMAliasOp(parser, result);
}
if (op === 'llvm.alloca') {
return this.parseLLVMAllocaOp(parser, result);
}
if (op === 'llvm.call') {
return this.parseLLVMCallOp(parser, result);
}
if (op === 'llvm.call_intrinsic') {
return this.parseLLVMCallIntrinsicOp(parser, result);
}
if (op === 'llvm.invoke') {
return this.parseLLVMInvokeOp(parser, result);
}
if (op === 'llvm.landingpad') {
return this.parseLLVMLandingpadOp(parser, result);
}
if (op === 'llvm.icmp' || op === 'llvm.fcmp') {
return this.parseLLVMCmpOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseLLVMGlobalOp(parser, result) {
const linkageKeywords = ['external', 'available_externally', 'linkonce', 'linkonce_odr', 'weak', 'weak_odr', 'appending', 'internal', 'private', 'extern_weak', 'common'];
const __linkage = parser.parseOptionalKeyword(linkageKeywords);
if (__linkage) {
result.addAttribute('linkage', __linkage);
}
const visibilityKeywords = ['default', 'hidden', 'protected'];
const __visibility = parser.parseOptionalKeyword(visibilityKeywords);
if (__visibility) {
result.addAttribute('visibility_', __visibility);
}
if (parser.parseOptionalKeyword('thread_local')) {
result.addAttribute('thread_local_', true);
}
const unnamedAddrKeywords = ['unnamed_addr', 'local_unnamed_addr'];
const __unnamedAddr = parser.parseOptionalKeyword(unnamedAddrKeywords);
if (__unnamedAddr) {
result.addAttribute('unnamed_addr', __unnamedAddr);
}
if (parser.parseOptionalKeyword('constant')) {
result.addAttribute('constant', true);
}
const _symName = parser.parseOptionalSymbolName();
if (_symName) {
result.addAttribute('sym_name', _symName);
}
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
const value = parser.parseAttribute();
if (parser.parseOptionalColon()) {
parser.parseType();
}
result.addAttribute('value', value);
parser.parseRParen();
}
if (parser.parseOptionalKeyword('comdat')) {
parser.parseLParen();
const comdat = parser.parseOptionalSymbolName();
parser.parseRParen();
result.addAttribute('comdat', comdat);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
}
parser.parseOptionalRegion(result.addRegion());
return true;
}
parseLLVMAliasOp(parser, result) {
const linkageKeywords = ['external', 'available_externally', 'linkonce', 'linkonce_odr', 'weak', 'weak_odr', 'internal', 'private'];
const __linkage = parser.parseOptionalKeyword(linkageKeywords);
if (__linkage) {
result.addAttribute('linkage', __linkage);
}
const visibilityKeywords = ['default', 'hidden', 'protected'];
const __visibility = parser.parseOptionalKeyword(visibilityKeywords);
if (__visibility) {
result.addAttribute('visibility_', __visibility);
}
if (parser.parseOptionalKeyword('thread_local')) {
result.addAttribute('thread_local_', true);
}
const unnamedAddrKeywords = ['unnamed_addr', 'local_unnamed_addr'];
const __unnamedAddr = parser.parseOptionalKeyword(unnamedAddrKeywords);
if (__unnamedAddr) {
result.addAttribute('unnamed_addr', __unnamedAddr);
}
const __aliasSym = parser.parseOptionalSymbolName();
if (__aliasSym) {
result.attributes.set('sym_name', __aliasSym);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addAttribute('alias_type', type);
}
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
parseGEPIndices(parser, op, operands, attrName) {
const rawConstantIndices = [];
do {
const constIndex = parser.parseOptionalInteger();
if (constIndex === null) {
const operand = parser.parseOperand();
operands.push(operand);
rawConstantIndices.push(-2147483648);
} else {
rawConstantIndices.push(constIndex);
}
} while (parser.parseOptionalComma());
if (rawConstantIndices.length > 0) {
op.addAttribute(attrName || 'rawConstantIndices', rawConstantIndices);
}
}
parseIndirectBrOpSucessors(parser, op /*, args */) {
// All operands listed first, then colon, then all types
parser.parseLSquare();
const segmentSizes = [];
if (!parser.parseOptionalRSquare()) {
do {
const successor = parser.parseOptionalSuccessor();
if (!op.successors) {
op.successors = [];
}
op.successors.push({ name: successor });
const unresolvedOperands = [];
const types = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
} while (parser.parseOptionalComma());
if (parser.parseOptionalColon()) {
do {
const type = parser.parseType();
types.push(type);
} while (parser.parseOptionalComma());
}
parser.parseRParen();
}
}
for (let i = 0; i < unresolvedOperands.length; i++) {
const type = i < types.length ? types[i] : null;
parser.resolveOperand(unresolvedOperands[i], type, op.operands);
}
segmentSizes.push(unresolvedOperands.length);
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (segmentSizes.length > 0) {
op.addAttribute('indbr_operand_segments', segmentSizes);
}
}
parseLLVMAllocaOp(parser, result) {
// llvm.alloca [inalloca] %arraySize x !elemType : (i64) -> !llvm.ptr
if (parser.parseOptionalKeyword('inalloca')) {
result.addAttribute('inalloca', true);
}
const arraySize = parser.parseOperand();
parser.parseKeyword('x');
const elemType = parser.parseType();
result.addAttribute('elem_type', elemType);
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const fnType = parser.parseType();
if (fnType instanceof _.FunctionType) {
parser.resolveOperands([arraySize], fnType.inputs, result.operands);
result.addTypes(fnType.results);
}
return true;
}
parseLLVMCallOp(parser, result) {
// llvm.call [cconv] [tailcall] @callee|%ptr (args) [vararg(type)] : func_type
const cconvKeywords = ['ccc', 'fastcc', 'coldcc', 'cc', 'webkit_jscc', 'anyregcc', 'preserve_mostcc', 'preserve_allcc', 'preserve_nonecc', 'cxx_fast_tlscc', 'tailcc', 'swiftcc', 'swifttailcc', 'cfguard_checkcc', 'ghccc', 'arm_apcscc', 'arm_aapcscc', 'arm_aapcs_vfpcc', 'aarch64_vector_pcs', 'aarch64_sve_vector_pcs', 'aarch64_sme_preservemost_from_x0', 'aarch64_sme_preservemost_from_x2', 'msp430_intrcc', 'avr_intrcc', 'avr_signalcc', 'ptx_kernelcc', 'ptx_devicecc', 'spir_funccc', 'spir_kernelcc', 'intel_ocl_bicc', 'x86_64_sysvcc', 'win64cc', 'x86_fastcallcc', 'x86_stdcallcc', 'x86_thiscallcc', 'x86_vectorcallcc', 'x86_intrcc', 'amdgpu_vs', 'amdgpu_gs', 'amdgpu_ps', 'amdgpu_cs', 'amdgpu_kernel', 'amdgpu_kernelcc', 'x86_regcallcc', 'amdgpu_hs', 'msp430_builtincc', 'amdgpu_ls', 'amdgpu_es', 'aarch64_vfpcc', 'aarch64_sve_vfpcc', 'wasm_emscripten_invokecc', 'amdgpu_gfx', 'm68k_intrcc', 'cc_10', 'cc_11'];
const __cconv = parser.parseOptionalKeyword(cconvKeywords);
if (__cconv) {
result.addAttribute('CConv', __cconv);
}
const tailcallKeywords = ['none', 'tail', 'musttail', 'notail'];
const __tailcall = parser.parseOptionalKeyword(tailcallKeywords);
if (__tailcall) {
result.addAttribute('TailCallKind', __tailcall);
}
let isDirect = false;
let calleePtr = null;
const callee = parser.parseOptionalSymbolName();
if (callee) {
result.addAttribute('callee', callee);
isDirect = true;
} else {
calleePtr = parser.parseOptionalOperand();
}
const unresolvedOperands = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
unresolvedOperands.push(arg);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
if (parser.parseOptionalKeyword('vararg')) {
parser.parseLParen();
const varCalleeType = parser.parseType();
result.addAttribute('var_callee_type', varCalleeType);
parser.parseRParen();
}
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
const opBundles = [];
do {
const tag = parser.parseString();
parser.parseLParen();
const bundleOperands = [];
const bundleTypes = [];
if (!parser.parseOptionalRParen()) {
do {
bundleOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseColon();
do {
bundleTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
opBundles.push({ tag, operands: bundleOperands, types: bundleTypes });
} while (parser.parseOptionalComma());
parser.parseRSquare();
if (opBundles.length > 0) {
result.addAttribute('op_bundle_tags', opBundles);
}
}
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
let calleePtrType = null;
if (!isDirect) {
calleePtrType = parser.parseType();
parser.parseComma();
}
const sig = parser.parseFunctionSignature();
if (calleePtr) {
parser.resolveOperand(calleePtr, calleePtrType, result.operands);
}
parser.resolveOperands(unresolvedOperands, sig.argTypes, result.operands);
result.addTypes(sig.resultTypes);
return true;
}
parseLLVMCallIntrinsicOp(parser, result) {
const intrinName = parser.parseString();
result.addAttribute('intrin', intrinName);
const unresolvedOperands = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const arg = parser.parseOperand();
unresolvedOperands.push(arg);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
const opBundles = [];
do {
const tag = parser.parseString();
parser.parseLParen();
const bundleOperands = [];
const bundleTypes = [];
if (!parser.parseOptionalRParen()) {
do {
bundleOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseColon();
do {
bundleTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
opBundles.push({ tag, operands: bundleOperands, types: bundleTypes });
} while (parser.parseOptionalComma());
parser.parseRSquare();
if (opBundles.length > 0) {
result.addAttribute('op_bundle_tags', opBundles);
}
}
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const sig = parser.parseFunctionSignature();
parser.resolveOperands(unresolvedOperands, sig.argTypes, result.operands);
result.addTypes(sig.resultTypes);
return true;
}
parseLLVMCmpOp(parser, result) {
// llvm.icmp "eq" %lhs, %rhs : i32
const predicate = parser.parseString();
result.addAttribute('predicate', predicate);
const lhs = parser.parseOperand();
parser.parseComma();
const rhs = parser.parseOperand();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const type = parser.parseType();
parser.resolveOperands([lhs, rhs], [type, type], result.operands);
// Result type is i1 (or vector/tensor of i1 for vector/tensor operands)
let resultType = new _.IntegerType('i1');
if (type instanceof _.VectorType) {
resultType = new _.VectorType(type.shape, new _.IntegerType('i1'), type.scalableDims);
} else if (type instanceof _.RankedTensorType) {
resultType = new _.RankedTensorType(type.shape, new _.IntegerType('i1'), type.encoding);
}
result.addTypes([resultType]);
return true;
}
parseLLVMInvokeOp(parser, result) {
const cconvKeywords = ['ccc', 'fastcc', 'coldcc', 'cc', 'webkit_jscc', 'anyregcc', 'preserve_mostcc', 'preserve_allcc', 'preserve_nonecc', 'cxx_fast_tlscc', 'tailcc', 'swiftcc', 'swifttailcc', 'cfguard_checkcc', 'ghccc', 'cc_10', 'cc_11'];
const __cconv = parser.parseOptionalKeyword(cconvKeywords);
if (__cconv) {
result.addAttribute('CConv', __cconv);
}
let isDirect = false;
let funcPtr = null;
const invokeCallee = parser.parseOptionalSymbolName();
if (invokeCallee) {
isDirect = true;
result.addAttribute('callee', invokeCallee);
} else {
funcPtr = parser.parseOptionalOperand();
}
const unresolvedOperands = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseKeyword('to');
const normalDest = parser.parseOptionalSuccessor();
result.successors = result.successors || [];
const normalSucc = { label: normalDest };
if (parser.parseOptionalLParen()) {
normalSucc.operands = [];
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
normalSucc.operands.push(operand);
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
result.successors.push(normalSucc);
parser.parseKeyword('unwind');
const unwindDest = parser.parseOptionalSuccessor();
const unwindSucc = { label: unwindDest };
if (parser.parseOptionalLParen()) {
unwindSucc.operands = [];
if (!parser.parseOptionalRParen()) {
do {
const operand = parser.parseOperand();
unwindSucc.operands.push(operand);
if (parser.parseOptionalColon()) {
parser.parseType();
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
result.successors.push(unwindSucc);
if (parser.parseOptionalKeyword('vararg')) {
parser.parseLParen();
const varargType = parser.parseType();
result.addAttribute('var_callee_type', varargType);
parser.parseRParen();
}
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
const opBundles = [];
do {
const tag = parser.parseString();
parser.parseLParen();
const bundleOperands = [];
const bundleTypes = [];
if (!parser.parseOptionalRParen()) {
do {
bundleOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseColon();
do {
bundleTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
opBundles.push({ tag, operands: bundleOperands, types: bundleTypes });
} while (parser.parseOptionalComma());
parser.parseRSquare();
if (opBundles.length > 0) {
result.addAttribute('op_bundle_tags', opBundles);
}
}
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
let calleePtrType = null;
if (!isDirect) {
calleePtrType = parser.parseType();
parser.parseComma();
}
const sig = parser.parseFunctionSignature();
if (funcPtr) {
parser.resolveOperand(funcPtr, calleePtrType, result.operands);
}
parser.resolveOperands(unresolvedOperands, sig.argTypes, result.operands);
result.addTypes(sig.resultTypes);
return true;
}
parseLLVMLandingpadOp(parser, result) {
if (parser.parseOptionalKeyword('cleanup')) {
result.addAttribute('cleanup', true);
}
while (parser.parseOptionalLParen()) {
parser.parseKeyword(); // 'catch' or 'filter'
const operand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(operand, type, result.operands);
parser.parseRParen();
}
parser.parseColon();
const resultType = parser.parseType();
result.addTypes([resultType]);
return true;
}
parseLLVMFuncOp(parser, result) {
const linkageKeywords = ['external', 'available_externally', 'linkonce', 'linkonce_odr', 'weak', 'weak_odr', 'appending', 'internal', 'private', 'extern_weak', 'common'];
const __linkage = parser.parseOptionalKeyword(linkageKeywords);
if (__linkage) {
result.addAttribute('linkage', __linkage);
}
const visibilityKeywords = ['default', 'hidden', 'protected'];
const __visibility = parser.parseOptionalKeyword(visibilityKeywords);
if (__visibility) {
result.addAttribute('visibility_', __visibility);
}
const unnamedAddrKeywords = ['unnamed_addr', 'local_unnamed_addr'];
const __unnamedAddr = parser.parseOptionalKeyword(unnamedAddrKeywords);
if (__unnamedAddr) {
result.addAttribute('unnamed_addr', __unnamedAddr);
}
const cconvKeywords = ['ccc', 'fastcc', 'coldcc', 'cc', 'webkit_jscc', 'anyregcc', 'preserve_mostcc', 'preserve_allcc', 'preserve_nonecc', 'cxx_fast_tlscc', 'tailcc', 'swiftcc', 'swifttailcc', 'cfguard_checkcc', 'ghccc', 'arm_apcscc', 'arm_aapcscc', 'arm_aapcs_vfpcc', 'aarch64_vector_pcs', 'aarch64_sve_vector_pcs', 'aarch64_sme_preservemost_from_x0', 'aarch64_sme_preservemost_from_x2', 'msp430_intrcc', 'avr_intrcc', 'avr_signalcc', 'ptx_kernelcc', 'ptx_devicecc', 'spir_funccc', 'spir_kernelcc', 'intel_ocl_bicc', 'x86_64_sysvcc', 'win64cc', 'x86_fastcallcc', 'x86_stdcallcc', 'x86_thiscallcc', 'x86_vectorcallcc', 'x86_intrcc', 'amdgpu_vs', 'amdgpu_gs', 'amdgpu_ps', 'amdgpu_cs', 'amdgpu_kernel', 'amdgpu_kernelcc', 'x86_regcallcc', 'amdgpu_hs', 'msp430_builtincc', 'amdgpu_ls', 'amdgpu_es', 'aarch64_vfpcc', 'aarch64_sve_vfpcc', 'wasm_emscripten_invokecc', 'amdgpu_gfx', 'm68k_intrcc', 'cc_10', 'cc_11'];
const __cconv = parser.parseOptionalKeyword(cconvKeywords);
if (__cconv) {
result.addAttribute('CConv', __cconv);
}
parser.parseSymbolName('sym_name', result.attributes);
const argResult = parser.parseFunctionArgumentList(true);
const params = argResult.arguments.map((a) => a.type);
const results = [];
const resultAttrs = [];
if (parser.parseOptionalArrow()) {
parser.parseFunctionResultList(results, resultAttrs);
}
const returnType = results.length > 0 ? results[0] : null;
const type = new _.llvm.LLVMFunctionType(returnType, params, argResult.isVariadic);
result.addAttribute('function_type', new _.TypeAttrOf(type));
if (parser.parseOptionalKeyword('vscale_range')) {
parser.parseLParen();
const minRange = parser.parseInteger();
parser.parseComma();
const maxRange = parser.parseInteger();
parser.parseRParen();
result.addAttribute('vscale_range', `(${minRange}, ${maxRange})`);
}
if (parser.parseOptionalKeyword('comdat')) {
parser.parseLParen();
const comdat = parser.parseOptionalSymbolName();
parser.parseRParen();
result.addAttribute('comdat', comdat);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion(), argResult.arguments, /* isIsolatedNameScope */ true);
return true;
}
};
_.ROCDLDialect = class extends _.llvm.LLVMDialect {
constructor(operations) {
super(operations, 'rocdl');
this.registerCustomAttribute('ROCDL_SchedGroupMaskAttr', this.parseEnumFlagsPipe.bind(this));
this.registerCustomAttribute('ROCDL_MFMANegModifierAttr', this.parseEnumFlagsPipe.bind(this));
this.registerCustomDirective('CachePolicy', this.parseCachePolicy.bind(this));
}
parseCachePolicy(parser, op, aux) {
const value = parser.parseOptionalInteger();
if (value !== null) {
op.addAttribute(aux, String(value));
return;
}
const family = parser.parseKeyword();
parser.parseLess();
const flags = [];
do {
flags.push(parser.parseKeyword());
} while (parser.parseOptionalVerticalBar());
parser.parseGreater();
op.addAttribute(aux, `${family}<${flags.join('|')}>`);
}
};
_.XSMMDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xsmm');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'xsmm.unary.invoke') {
return this.parseUnaryInvokeOp(parser, result);
}
if (op === 'xsmm.gemm.invoke.f32') {
return this.parseGemmInvokeF32Op(parser, result);
}
if (op === 'xsmm.gemm.relu.invoke.f32' || op === 'xsmm.gemm.bias.invoke.f32' || op === 'xsmm.gemm.relu.beta.invoke.f32') {
return this.parseGemm4InvokeF32Op(parser, result);
}
if (op === 'xsmm.brgemm.offs.invoke.f32') {
return this.parseBRGemmOffsInvokeF32Op(parser, result);
}
if (op === 'xsmm.brgemm.offs.relu.invoke.f32' || op === 'xsmm.brgemm.offs.bias.invoke.f32' || op === 'xsmm.brgemm.offs.relu.beta1.invoke.f32') {
return this.parseBRGemmOffs4InvokeF32Op(parser, result);
}
return super.parseOperation(parser, result);
}
parseUnaryInvokeOp(parser, result) {
const output = parser.parseOperand();
const outputIndices = parser.parseOperandList('square');
parser.parseEqual();
const ptr = parser.parseOperand();
parser.parseLParen();
const input = parser.parseOperand();
const inputIndices = parser.parseOperandList('square');
parser.parseRParen();
parser.parseColon();
parser.parseType();
for (const op of [ptr, output, input, ...outputIndices, ...inputIndices]) {
parser.resolveOperand(op, null, result.operands);
}
return true;
}
parseGemmInvokeF32Op(parser, result) {
const ptr = parser.parseOperand();
parser.parseComma();
const c = parser.parseOperand();
const cIndices = parser.parseOperandList('square');
parser.parseEqual();
const a = parser.parseOperand();
const aIndices = parser.parseOperandList('square');
parser.parseComma();
const b = parser.parseOperand();
const bIndices = parser.parseOperandList('square');
parser.parseColon();
parser.parseType();
for (const op of [ptr, c, a, b, ...cIndices, ...aIndices, ...bIndices]) {
parser.resolveOperand(op, null, result.operands);
}
return true;
}
parseGemm4InvokeF32Op(parser, result) {
const ptr = parser.parseOperand();
parser.parseComma();
const c = parser.parseOperand();
const cIndices = parser.parseOperandList('square');
parser.parseEqual();
const a = parser.parseOperand();
const aIndices = parser.parseOperandList('square');
parser.parseComma();
const b = parser.parseOperand();
const bIndices = parser.parseOperandList('square');
parser.parseComma();
const d = parser.parseOperand();
const dIndices = parser.parseOperandList('square');
parser.parseColon();
parser.parseType();
for (const op of [ptr, c, a, b, d, ...cIndices, ...aIndices, ...bIndices, ...dIndices]) {
parser.resolveOperand(op, null, result.operands);
}
return true;
}
parseBRGemmOffsInvokeF32Op(parser, result) {
const ptr = parser.parseOperand();
parser.parseComma();
const c = parser.parseOperand();
const cIndices = parser.parseOperandList('square');
parser.parseEqual();
const a = parser.parseOperand();
const aIndices = parser.parseOperandList('square');
parser.parseComma();
const b = parser.parseOperand();
const bIndices = parser.parseOperandList('square');
parser.parseComma();
parser.parseKeyword('aOffsets');
parser.parseEqual();
result.addAttribute('aOffsets', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseComma();
parser.parseKeyword('bOffsets');
parser.parseEqual();
result.addAttribute('bOffsets', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseComma();
parser.parseKeyword('numBatches');
parser.parseEqual();
result.addAttribute('numBatches', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseColon();
parser.parseType();
for (const op of [ptr, c, a, b, ...cIndices, ...aIndices, ...bIndices]) {
parser.resolveOperand(op, null, result.operands);
}
return true;
}
parseBRGemmOffs4InvokeF32Op(parser, result) {
const ptr = parser.parseOperand();
parser.parseComma();
const c = parser.parseOperand();
const cIndices = parser.parseOperandList('square');
parser.parseEqual();
const a = parser.parseOperand();
const aIndices = parser.parseOperandList('square');
parser.parseComma();
const b = parser.parseOperand();
const bIndices = parser.parseOperandList('square');
parser.parseComma();
const d = parser.parseOperand();
const dIndices = parser.parseOperandList('square');
parser.parseComma();
parser.parseKeyword('aOffsets');
parser.parseEqual();
result.addAttribute('aOffsets', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseComma();
parser.parseKeyword('bOffsets');
parser.parseEqual();
result.addAttribute('bOffsets', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseComma();
parser.parseKeyword('numBatches');
parser.parseEqual();
result.addAttribute('numBatches', parser.parseAttribute(new _.IntegerType('i64')));
parser.parseColon();
parser.parseType();
for (const op of [ptr, c, a, b, d, ...cIndices, ...aIndices, ...bIndices, ...dIndices]) {
parser.resolveOperand(op, null, result.operands);
}
return true;
}
};
_.StdxDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'stdx');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'stdx.closure') {
return this.parseClosureOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseClosureOp(parser, result) {
const sig = parser.parseFunctionSignatureWithArguments(false);
const argTypes = sig.arguments.map((a) => a.type);
const type = { inputs: argTypes, results: sig.resultTypes };
result.addAttribute('type', type);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region, sig.arguments);
return true;
}
};
_.vm = {};
_.vm.RefType = class extends _.Type {
constructor(objectType) {
super(null);
this.objectType = objectType;
}
static parse(parser) {
parser.parseLess();
const objectType = parser.parseOptionalQuestion()
? new _.OpaqueType('?')
: parser.parseType();
parser.parseGreater();
return new _.vm.RefType(objectType);
}
toString() {
return `!vm.ref<${this.objectType}>`;
}
};
_.vm.VMDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'vm');
this.registerCustomDirective('BranchTableCases', this.parseBranchTableCases.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'vm.func') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'vm.cond_fail') {
// or: vm.cond_fail %status, "message"
// or: vm.cond_fail %cond, %status
// or: vm.cond_fail %status
const unresolvedOperands = [];
const firstOp = parser.parseOperand();
unresolvedOperands.push(firstOp);
if (parser.parseOptionalComma()) {
// Could be second operand or message
const __secondOp = parser.parseOptionalOperand();
if (__secondOp) {
unresolvedOperands.push(__secondOp);
// Optional message
if (parser.parseOptionalComma()) {
const msg = parser.parseOptionalString();
if (msg !== null) {
result.addAttribute('message', msg);
}
}
} else {
const msg = parser.parseOptionalString();
if (msg !== null) {
result.addAttribute('message', msg);
}
}
}
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, result.operands);
}
return true;
}
if (op === 'vm.import') {
parser.parseOptionalVisibilityKeyword(result.attributes);
if (parser.parseOptionalKeyword('optional')) {
result.addAttribute('is_optional', true);
}
parser.parseSymbolName('sym_name', result.attributes);
const inputs = [];
parser.parseLParen();
while (!parser.parseOptionalRParen()) {
if (parser.parseOptionalOperand()) {
inputs.push(parser.parseColonType());
} else {
inputs.push(parser.parseType());
}
parser.parseOptionalEllipsis();
parser.parseOptionalComma();
}
const results = [];
const resultAttrs = [];
if (parser.parseOptionalArrow()) {
parser.parseFunctionResultList(results, resultAttrs);
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'vm.export') {
const functionRef = parser.parseOptionalSymbolName();
result.addAttribute('function_ref', functionRef);
if (parser.parseOptionalKeyword('as')) {
parser.parseLParen();
const exportName = parser.parseString();
result.addAttribute('export_name', exportName);
parser.parseRParen();
} else {
result.addAttribute('export_name', functionRef);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
return true;
}
if (op === 'vm.initializer') {
parser.parseOptionalVisibilityKeyword(result.attributes);
const __symName = parser.parseOptionalSymbolName();
if (__symName !== null) {
result.attributes.set('sym_name', __symName);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'vm.const.i32.zero') {
result.compatibility = true;
const __symbol = parser.parseOptionalSymbolName();
if (__symbol === null) {
const value = parser.parseOptionalAttribute();
if (value !== null) {
result.addAttribute('value', value.value === undefined ? value : value.value);
}
} else {
result.addAttribute('rodata', __symbol);
}
parser.parseOptionalAttrDict(result.attributes);
const types = parser.parseOptionalColonTypeList();
result.addTypes(types.length > 0 ? types : [new _.IntegerType('i32')]);
return true;
}
// Handle vm.switch.ref operation
if (op === 'vm.switch.ref') {
const unresolvedOperands = [];
const indexUnresolved = parser.parseOperand();
unresolvedOperands.push(indexUnresolved);
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
const __val = parser.parseOptionalOperand();
if (__val) {
unresolvedOperands.push(__val);
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parser.parseKeyword('else');
const defaultValueUnresolved = parser.parseOperand();
unresolvedOperands.push(defaultValueUnresolved);
parser.parseOptionalAttrDict(result.attributes);
let resultType = null;
if (parser.parseOptionalColon()) {
resultType = parser.parseType();
result.addTypes([resultType]);
}
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, resultType, result.operands);
}
return true;
}
// Following reference: CallVariadicOp::parse in VMOps.cpp:1251
if (op === 'vm.call.variadic') {
const calleeAttr = parser.parseAttribute();
result.addAttribute('callee', calleeAttr);
parser.parseLParen();
const flatOperands = [];
const flatSegmentSizes = [];
while (!parser.parseOptionalRParen()) {
if (parser.parseOptionalLSquare()) {
// Variadic list
const flatSegmentOperands = [];
while (!parser.parseOptionalRSquare()) {
if (parser.parseOptionalLParen()) {
// Tuple inside variadic list
while (!parser.parseOptionalRParen()) {
flatSegmentOperands.push(parser.parseOperand());
if (!parser.parseOptionalComma()) {
parser.parseRParen();
break;
}
}
} else {
flatSegmentOperands.push(parser.parseOperand());
}
if (!parser.parseOptionalComma()) {
parser.parseRSquare();
break;
}
}
flatSegmentSizes.push(flatSegmentOperands.length);
flatOperands.push(...flatSegmentOperands);
} else {
// Normal single operand
flatOperands.push(parser.parseOperand());
flatSegmentSizes.push(-1);
}
if (!parser.parseOptionalComma()) {
parser.parseRParen();
break;
}
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
parser.parseLParen();
const flatOperandTypes = [];
const segmentTypes = [];
const segmentSizes = [];
let segmentIndex = 0;
while (!parser.parseOptionalRParen()) {
const operandType = parser.parseType();
const isVariadic = parser.parseOptionalEllipsis();
if (isVariadic) {
const flatSegmentSize = flatSegmentSizes[segmentIndex];
for (let i = 0; i < flatSegmentSize; i++) {
flatOperandTypes.push(operandType);
}
segmentSizes.push(flatSegmentSize);
} else {
flatOperandTypes.push(operandType);
segmentSizes.push(-1);
}
segmentTypes.push(operandType);
segmentIndex++;
if (!parser.parseOptionalComma()) {
parser.parseRParen();
break;
}
}
parser.resolveOperands(flatOperands, flatOperandTypes, result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
return super.parseOperation(parser, result);
}
parseBranchTableCases(parser, op /*, args */) {
if (parser.parseOptionalKeyword('default')) {
parser.parseColon();
const defaultDest = parser.parseOptionalSuccessor();
op.successors = op.successors || [];
const succ = { dest: defaultDest };
if (parser.parseOptionalLParen()) {
const operands = parser.parseOperandList('none');
const types = parser.parseOptionalColonTypeList();
parser.parseRParen();
parser.resolveOperands(operands, types);
succ.operands = operands;
}
op.successors.push(succ);
parser.parseOptionalComma();
}
const caseValues = [];
for (;;) {
const caseValue = parser.parseOptionalInteger();
if (caseValue === null) {
break;
}
caseValues.push(caseValue);
parser.parseColon();
const caseDest = parser.parseOptionalSuccessor();
const caseSucc = { dest: caseDest };
if (parser.parseOptionalLParen()) {
const operands = parser.parseOperandList('none');
const types = parser.parseOptionalColonTypeList();
parser.parseRParen();
parser.resolveOperands(operands, types);
caseSucc.operands = operands;
}
op.successors.push(caseSucc);
parser.parseOptionalComma();
}
if (caseValues.length > 0) {
op.addAttribute('case_values', caseValues);
}
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'ref') {
return _.vm.RefType.parse(parser);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
};
_.MathDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'math');
this.registerCustomAttribute('Arith_FastMathAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
}
};
_.torch.TMTensorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tm_tensor');
}
};
_.MLProgramDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'ml_program');
this.registerCustomDirective('TypedInitialValue', this.parseTypedInitialValue.bind(this));
this.registerCustomDirective('TokenOrdering', this.parseTokenOrdering.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'ml_program.func' || op === 'ml_program.subgraph') {
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
parseTokenOrdering(parser, op, consumeTokens, produceTokenType) {
if (!parser.parseOptionalKeyword('ordering')) {
return;
}
parser.parseLParen();
if (parser.parseOptionalLParen()) {
parser.parseRParen();
} else {
let __tok = parser.parseOptionalOperand();
while (__tok) {
consumeTokens.push(__tok);
if (!parser.parseOptionalComma()) {
break;
}
__tok = parser.parseOptionalOperand();
}
}
parser.parseArrow();
produceTokenType.push(parser.parseType());
parser.parseRParen();
}
parseTypedInitialValue(parser, op, typeAttr, valueAttr) {
if (parser.parseOptionalLParen()) {
const attr = parser.parseAttribute();
if (parser.parseOptionalColon()) {
attr.type = parser.parseType();
}
parser.parseRParen();
op.addAttribute(valueAttr, attr.value === undefined ? attr : attr.value);
}
parser.parseColon();
const type = parser.parseType();
op.addAttribute(typeAttr, type);
}
};
_.IREEGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'iree_gpu');
this.registerCustomDirective('AsyncDMASourceIndexTypes', this.parseAsyncDMASourceIndexTypes.bind(this));
}
parseAsyncDMASourceIndexTypes(parser, op, sourceIndicesOps, sourceIndicesTypes) {
if (parser.parseOptionalLSquare()) {
for (const t of parser.parseTypeList()) {
sourceIndicesTypes.push(t);
}
parser.parseRSquare();
} else {
const indexType = new _.IndexType();
for (let i = 0; i < sourceIndicesOps.length; i++) {
sourceIndicesTypes.push(indexType);
}
}
}
};
_.TFDeviceDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tf_device');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tf_device.replicate') {
return this.parseReplicateOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseReplicateOp(parser, result) {
let n = 1;
if (!parser.parseOptionalLParen()) {
parser.parseOptionalAttrDict(result.attributes);
} else if (parser.parseOptionalRParen()) {
parser.parseOptionalAttrDict(result.attributes);
} else {
do {
if (parser.parseOptionalLSquare()) {
const unresolvedInputs = [];
while (!parser.parseOptionalRSquare()) {
unresolvedInputs.push(parser.parseOperand());
if (!parser.parseOptionalComma()) {
parser.parseRSquare();
break;
}
}
parser.parseKeyword('as');
parser.parseOperand(); // block arg
parser.parseColon();
const type = parser.parseType();
for (const input of unresolvedInputs) {
parser.resolveOperand(input, type, result.operands);
}
} else {
const unresolvedValue = parser.parseOptionalOperand();
if (!unresolvedValue) {
break;
}
parser.parseKeyword('as');
parser.parseOperand(); // block arg
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(unresolvedValue, type, result.operands);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
parser.parseOptionalAttrDict(result.attributes);
}
n = result.attributes.get('n').value;
{
const region = result.addRegion();
parser.parseRegion(region);
if (region.blocks.length > 0) {
const block = region.blocks[0];
if (block.operations.length > 0) {
const terminator = block.operations[block.operations.length - 1];
if (terminator.operands) {
for (const operand of terminator.operands) {
if (operand.type) {
for (let i = 0; i < n; i++) {
result.addTypes([operand.type]);
}
}
}
}
}
}
}
return true;
}
};
_.TFGDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfg');
}
getOperation(opName) {
let op = super.getOperation(opName);
if (!op) {
this.registerOperandName(opName, {});
op = super.getOperation(opName);
}
return op;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tfg.func') {
if (parser.parseOptionalKeyword('generic')) {
result.addAttribute('generic', true);
}
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'tfg.return') {
let dataOperands = [];
if (parser.parseOptionalLParen()) {
dataOperands = parser.parseOperandList('none');
parser.parseRParen();
}
const controlOperands = [];
const controlRetAttrs = [];
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const __ctlDep = parser.parseOptionalOperand();
if (__ctlDep) {
controlOperands.push(__ctlDep);
const __ctlAttrs = new Map();
if (parser.parseOptionalAttrDict(__ctlAttrs)) {
controlRetAttrs.push(Object.fromEntries(__ctlAttrs));
} else {
controlRetAttrs.push({});
}
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
if (controlRetAttrs.length > 0) {
result.addAttribute('control_ret_attrs', controlRetAttrs);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const types = parser.parseTypeList();
parser.resolveOperands(dataOperands, types, result.operands);
} else {
parser.resolveOperands(dataOperands, dataOperands.map(() => null), result.operands);
}
parser.resolveOperands(controlOperands, controlOperands.map(() => null), result.operands);
return true;
}
const opInfo = result.name.getRegisteredInfo();
if (!opInfo.metadata.assemblyFormat) {
this.parseTFGOperation(parser, result);
return true;
}
return super.parseOperation(parser, result);
}
parseTFGOperation(parser, result) {
let unresolvedArgs = [];
if (parser.parseOptionalLParen()) {
unresolvedArgs = parser.parseOperandList('none');
parser.parseRParen();
}
const unresolvedCtls = [];
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
const operand = parser.parseOptionalOperand();
if (operand) {
unresolvedCtls.push(operand);
}
if (!parser.parseOptionalComma()) {
parser.parseRSquare();
break;
}
}
}
if (parser.parseOptionalKeyword('device')) {
parser.parseLParen();
const device = parser.parseString();
parser.parseRParen();
result.addAttribute('device', device);
}
if (parser.parseOptionalKeyword('name')) {
parser.parseLParen();
const name = parser.parseString();
parser.parseRParen();
result.addAttribute('_mlir_name', name);
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedArgs, type.inputs, result.operands);
parser.resolveOperands(unresolvedCtls, unresolvedCtls.map(() => new _.Type('!tfg.control')), result.operands);
result.addTypes(type.results);
} else {
const types = [type];
while (parser.parseOptionalComma()) {
types.push(parser.parseType());
}
parser.resolveOperands(unresolvedArgs, types, result.operands);
parser.resolveOperands(unresolvedCtls, unresolvedCtls.map(() => new _.Type('!tfg.control')), result.operands);
}
} else {
parser.resolveOperands(unresolvedArgs, unresolvedArgs.map(() => null), result.operands);
parser.resolveOperands(unresolvedCtls, unresolvedCtls.map(() => new _.Type('!tfg.control')), result.operands);
}
result.addTypes([new _.Type('!tfg.control')]);
}
};
_.TFExecutorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tf_executor');
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'control' || mnemonic === 'token') {
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tf_executor.graph') {
return this.parseGraphOp(parser, result);
}
if (op === 'tf_executor.island') {
return this.parseIslandOp(parser, result);
}
if (op === 'tf_executor.Enter') {
return this.parseEnterOp(parser, result);
}
if (op === 'tf_executor._SwitchN') {
const unresolvedData = parser.parseOperand();
parser.parseComma();
const unresolvedIndex = parser.parseOperand();
parser.parseKeyword('of');
const numOuts = parser.parseInteger();
result.addAttribute('num_outs', numOuts);
let unresolvedControlInputs = [];
if (parser.parseOptionalLParen()) {
unresolvedControlInputs = parser.parseOperandList('none');
parser.parseRParen();
}
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(unresolvedData, type, result.operands);
parser.resolveOperand(unresolvedIndex, new _.RankedTensorType([], new _.IntegerType('i32'), null), result.operands);
parser.resolveOperands(unresolvedControlInputs, unresolvedControlInputs.map(() => new _.Type('!tf_executor.control')), result.operands);
for (let i = 0; i < numOuts; i++) {
result.addTypes([type]);
}
result.addTypes([new _.Type('!tf_executor.control')]);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'tf_executor.Switch' || op === 'tf_executor.Merge' ||
op === 'tf_executor.LoopCond' || op === 'tf_executor.Exit') {
// These ops have hasCustomAssemblyFormat: true but no assemblyFormat in metadata
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedOperands, type.inputs, result.operands);
result.addTypes(type.results);
} else {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => type), result.operands);
result.addTypes([type]);
if (op === 'tf_executor.Switch') {
result.addTypes([type]);
}
if (op === 'tf_executor.Merge') {
result.addTypes([new _.RankedTensorType([], new _.IntegerType('i32'), null)]);
}
result.addTypes([new _.Type('!tf_executor.control')]);
}
} else {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
return super.parseOperation(parser, result);
}
parseEnterOp(parser, result) {
const unresolvedOperands = [];
let operand = parser.parseOptionalOperand();
while (operand) {
unresolvedOperands.push(operand);
if (!parser.parseOptionalComma()) {
break;
}
operand = parser.parseOptionalOperand();
}
parser.parseKeyword('frame');
const frameName = parser.parseString();
result.addAttribute('frame_name', frameName);
if (parser.parseOptionalKeyword('parallel_iterations')) {
const parallelIterations = parser.parseInteger();
result.addAttribute('parallel_iterations', parseInt(parallelIterations, 10));
} else {
result.addAttribute('parallel_iterations', 10);
}
const isConstant = parser.parseOptionalKeyword('constant');
result.addAttribute('is_constant', isConstant);
parser.parseColon();
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedOperands, type.inputs, result.operands);
result.addTypes(type.results);
} else {
const resolveTypes = unresolvedOperands.map((v, i) => i === 0 ? type : new _.Type('!tf_executor.control'));
parser.resolveOperands(unresolvedOperands, resolveTypes, result.operands);
result.addTypes([type]);
result.addTypes([new _.Type('!tf_executor.control')]);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseGraphOp(parser, result) {
const region = result.addRegion();
if (parser.parseOptionalRegion(region)) {
if (region.blocks && region.blocks.length > 0) {
const [block] = region.blocks;
if (block.operations && block.operations.length > 0) {
const lastOp = block.operations[block.operations.length - 1];
if (lastOp.name.getStringRef() === 'tf_executor.fetch' && lastOp.operands) {
for (const operand of lastOp.operands) {
const typeStr = operand.type ? operand.type.toString() : '';
if (operand.type && typeStr !== '!tf_executor.control') {
result.addTypes([operand.type]);
}
}
}
}
}
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
parseIslandOp(parser, result) {
// or: tf_executor.island {...}
// or: tf_executor.island(%control_inputs) {...}
if (parser.parseOptionalLParen()) {
const unresolvedOperands = parser.parseOperandList('none');
parser.parseRParen();
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
const region = result.addRegion();
if (parser.parseOptionalKeyword('wraps')) {
const wrappedOp = parser.parseGenericOperation();
result.addAttribute('wrappedOp', wrappedOp);
for (const opResult of wrappedOp.results) {
result.addTypes([opResult.type]);
}
} else if (parser.parseOptionalRegion(region)) {
if (region.blocks.length > 0) {
const block = region.blocks[region.blocks.length - 1];
if (block.operations.length > 0) {
const terminator = block.operations[block.operations.length - 1];
if (terminator.name.getStringRef() === 'tf_executor.yield') {
for (const operand of terminator.operands) {
result.addTypes([operand.type]);
}
}
}
}
}
result.addTypes([new _.Type('!tf_executor.control')]);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
};
_.TFFrameworkDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tf_framework');
}
};
_.TFRDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfr');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tfr.func') {
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
};
_.CoreRTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'corert');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'corert.executeop' || op === 'corert.executeop.seq') {
const isSeq = op === 'corert.executeop.seq';
const opHandlerOperands = parser.parseOperandList('paren');
for (const operand of opHandlerOperands) {
parser.resolveOperand(operand, null, result.operands);
}
const opNameAttr = parser.parseString();
result.addAttribute('op_name', opNameAttr);
const operandOperands = parser.parseOperandList('paren');
for (const operand of operandOperands) {
parser.resolveOperand(operand, null, result.operands);
}
parser.parseOptionalAttrDict(result.attributes);
const funcAttrs = new Map();
parser.parseOptionalAttrDict(funcAttrs);
if (funcAttrs.size > 0) {
result.addAttribute('op_func_attrs', Object.fromEntries(funcAttrs));
}
if (parser.parseOptionalColon()) {
const resultCount = parser.parseInteger();
if (isSeq) {
result.addTypes([new _.Type('!tfrt.chain')]);
}
const tensorHandleType = new _.Type('!corert.tensorhandle');
for (let i = 0; i < resultCount; i++) {
result.addTypes([tensorHandleType]);
}
} else if (isSeq) {
result.addTypes([new _.Type('!tfrt.chain')]);
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.TFRTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfrt');
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
const opInfo = result.name.getRegisteredInfo();
if (!opInfo) {
return false;
}
if (opInfo.metadata.assemblyFormat === 'operands attr-dict') {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, result.operands);
}
this.inferResultTypes(result, new Map());
return true;
}
if (op === 'tfrt.call') {
parser.parseSymbolName('callee', result.attributes);
const unresolvedOperands = parser.parseOperandList('paren');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type) {
if (type.inputs) {
parser.resolveOperands(unresolvedOperands, type.inputs, result.operands);
}
if (type.results) {
type.results.forEach((resultType) => {
result.addTypes([resultType]);
});
}
}
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
return true;
}
if (op === 'tfrt.return') {
const unresolvedOperands = parser.parseOperandList();
if (unresolvedOperands.length > 0) {
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
}
return true;
}
if (op === 'tfrt.repeat.i32') {
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
const types = parser.parseOptionalColonTypeList();
result.addTypes(types);
if (unresolvedOperands.length > 0) {
const i32Type = new _.IntegerType('i32');
parser.resolveOperand(unresolvedOperands[0], i32Type, result.operands);
const loopOperands = unresolvedOperands.slice(1);
parser.resolveOperands(loopOperands, types, result.operands);
}
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'tfrt.if') {
const unresolvedOperands = parser.parseOperandList();
if (parser.parseOptionalKeyword('attributes')) {
parser.parseOptionalAttrDict(result.attributes);
}
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType) {
if (funcType.inputs) {
parser.resolveOperands(unresolvedOperands, funcType.inputs, result.operands);
}
if (funcType.results) {
for (const resultType of funcType.results) {
result.addTypes([resultType]);
}
}
}
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
const thenRegion = {};
if (parser.parseOptionalRegion(thenRegion)) {
result.regions.push(thenRegion);
}
if (parser.parseOptionalKeyword('else')) {
const elseRegion = {};
if (parser.parseOptionalRegion(elseRegion)) {
result.regions.push(elseRegion);
}
}
return true;
}
if (op === 'tfrt.parallel_for.i32') {
const startUnresolved = parser.parseOperand();
parser.parseKeyword('to');
const endUnresolved = parser.parseOperand();
parser.parseKeyword('fixed');
const blockSizeUnresolved = parser.parseOperand();
let additionalArgs = [];
if (parser.parseOptionalComma()) {
additionalArgs = parser.parseOperandList();
}
const types = parser.parseOptionalColonTypeList();
const i32Type = new _.IntegerType('i32');
parser.resolveOperand(startUnresolved, i32Type, result.operands);
parser.resolveOperand(endUnresolved, i32Type, result.operands);
parser.resolveOperand(blockSizeUnresolved, i32Type, result.operands);
parser.resolveOperands(additionalArgs, types, result.operands);
result.addTypes([new _.Type('!tfrt.chain')]);
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'tfrt.parallel_call.i32') {
const startUnresolved = parser.parseOperand();
parser.parseKeyword('to');
const endUnresolved = parser.parseOperand();
parser.parseKeyword('fixed');
const blockSizeUnresolved = parser.parseOperand();
const callee = parser.parseOptionalSymbolName();
result.addAttribute('callee', callee);
const additionalArgs = parser.parseOperandList('paren');
const types = parser.parseOptionalColonTypeList();
const i32Type = new _.IntegerType('i32');
parser.resolveOperand(startUnresolved, i32Type, result.operands);
parser.resolveOperand(endUnresolved, i32Type, result.operands);
parser.resolveOperand(blockSizeUnresolved, i32Type, result.operands);
parser.resolveOperands(additionalArgs, types, result.operands);
result.addTypes([new _.Type('!tfrt.chain')]);
return true;
}
if (op === 'tfrt.while') {
const condUnresolved = parser.parseOperand();
const bodyFn = parser.parseOptionalSymbolName();
result.addAttribute('body_fn', bodyFn);
const argsUnresolved = parser.parseOperandList('paren');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalKeyword('parallel_iterations')) {
parser.parseLParen();
const parallelIterations = parser.parseInteger();
result.addAttribute('parallel_iterations', parseInt(parallelIterations, 10));
parser.parseRParen();
}
parser.parseColon();
parser.parseLParen();
const inputTypes = parser.parseTypeList();
parser.parseRParen();
parser.parseArrow();
parser.parseLParen();
const resultTypes = parser.parseTypeList();
parser.parseRParen();
parser.resolveOperand(condUnresolved, new _.IntegerType('i1'), result.operands);
parser.resolveOperands(argsUnresolved, inputTypes, result.operands);
for (const resultType of resultTypes) {
result.addTypes([resultType]);
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.TFRTFallbackAsyncDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfrt_fallback_async');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
const opInfo = result.name.getRegisteredInfo();
if (!opInfo) {
return false;
}
if (op === 'tfrt_fallback_async.batch_function') {
this.parseExecuteOpCommon(parser, result, { has_device: true, has_symbol_ref: true });
return true;
}
if (op === 'tfrt_fallback_async.createop') {
this.parseExecuteOpCommon(parser, result, { has_chain: true, has_key: true, has_device: true, has_func_attr: true });
parser.parseKeyword('num_args');
parser.parseLParen();
const numArgs = parser.parseInteger();
parser.parseRParen();
result.addAttribute('num_args', parseInt(numArgs, 10));
return true;
}
if (op === 'tfrt_fallback_async.executeop') {
this.parseExecuteOpCommon(parser, result, { has_key: true, has_cost: true, has_device: true, has_func_attr: true });
return true;
}
if (op === 'tfrt_fallback_async.executeop.seq') {
this.parseExecuteOpCommon(parser, result, { has_chain: true, has_key: true, has_cost: true, has_device: true, has_func_attr: true });
return true;
}
if (op === 'tfrt_fallback_async.executeop.allocator') {
const allocator = parser.parseOperandList('paren');
for (const operand of allocator) {
parser.resolveOperand(operand, null, result.operands);
}
this.parseExecuteOpCommon(parser, result, { has_key: true, has_cost: true, has_device: true, has_func_attr: true });
return true;
}
if (op === 'tfrt_fallback_async.executeop.seq.allocator') {
const chainAndAllocator = parser.parseOperandList('paren');
for (const operand of chainAndAllocator) {
parser.resolveOperand(operand, null, result.operands);
}
result.addTypes([new _.Type('!tfrt.chain')]);
this.parseExecuteOpCommon(parser, result, { has_key: true, has_cost: true, has_device: true, has_func_attr: true });
return true;
}
return super.parseOperation(parser, result);
}
parseExecuteOpCommon(parser, result, options) {
if (options.has_chain) {
const chainOperands = parser.parseOperandList('paren');
for (const operand of chainOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
if (options.has_key) {
parser.parseKeyword('key');
parser.parseLParen();
const opKey = parser.parseAttribute();
parser.parseRParen();
result.addAttribute('op_key', opKey);
}
if (options.has_cost) {
parser.parseKeyword('cost');
parser.parseLParen();
const cost = parser.parseAttribute();
parser.parseRParen();
result.addAttribute('_tfrt_cost', cost);
}
if (options.has_device) {
parser.parseKeyword('device');
parser.parseLParen();
const device = parser.parseAttribute();
parser.parseRParen();
result.addAttribute('device', device);
}
if (options.has_op_name !== false) {
const opName = parser.parseOptionalString();
if (opName !== null) {
result.addAttribute('op_name', opName);
}
}
if (options.has_symbol_ref) {
const f = parser.parseOptionalSymbolName();
if (f) {
result.addAttribute('f', f);
}
}
const unresolvedOperands = parser.parseOperandList('paren');
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
parser.parseOptionalAttrDict(result.attributes);
if (options.has_func_attr) {
parser.parseOptionalAttrDict(result.attributes);
}
if (parser.parseOptionalColon()) {
const resultCount = parser.parseInteger();
if (options.has_chain) {
result.addTypes([new _.Type('!tfrt.chain')]);
}
for (let i = 0; i < resultCount; i++) {
result.addTypes([new _.Type('!tfrt_fallback.tf_tensor')]);
}
} else if (options.has_chain) {
result.addTypes([new _.Type('!tfrt.chain')]);
}
}
};
_.TileDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tile');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// tile.contract has format: tile.contract agg, combo, operands... attributes : types -> result
// Example: %1 = tile.contract add, mul, %0, %arg0, %arg1 {sink = #map0, srcs = [#map1, #map2]} : tensor<f32>, tensor<1x256xf32>, tensor<256x512xf32> -> tensor<1x512xf32>
if (op === 'tile.contract') {
const __agg = parser.parseOptionalKeyword();
if (__agg) {
result.addAttribute('agg', __agg);
}
parser.parseOptionalComma();
const __combo = parser.parseOptionalKeyword();
if (__combo) {
result.addAttribute('combo', __combo);
}
parser.parseOptionalComma();
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
parser.resolveOperands(unresolvedOperands, parser.parseOptionalColonTypeList(), result.operands);
result.addTypes(parser.parseOptionalArrowTypeList());
return true;
}
return super.parseOperation(parser, result);
}
};
_.PXADialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'pxa');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'pxa.reduce' || op === 'pxa.vector_reduce') {
const agg = parser.parseKeyword();
result.addAttribute('agg', agg);
const unresolvedVal = parser.parseOperand();
parser.parseOptionalComma();
const unresolvedMemref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
let memrefType = null;
let valType = null;
if (parser.parseOptionalColon()) {
memrefType = parser.parseType();
result.addTypes([memrefType]);
if (op === 'pxa.vector_reduce' && parser.parseOptionalComma()) {
valType = parser.parseType();
}
}
parser.resolveOperand(unresolvedVal, valType, result.operands);
parser.resolveOperand(unresolvedMemref, memrefType, result.operands);
return true;
}
if (op === 'pxa.load' || op === 'pxa.vector_load') {
const unresolvedMemref = parser.parseOperand();
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, 'map', result.attributes);
parser.parseOptionalAttrDict(result.attributes);
let memrefType = null;
if (parser.parseOptionalColon()) {
memrefType = parser.parseType();
if (op === 'pxa.vector_load' && parser.parseOptionalComma()) {
const vectorType = parser.parseType();
result.addTypes([vectorType]);
} else if (op === 'pxa.load' && memrefType) {
// Result type is element type of memref
const elementType = memrefType.elementType || memrefType;
result.addTypes([elementType]);
}
}
parser.resolveOperand(unresolvedMemref, memrefType, result.operands);
return true;
}
if (op === 'pxa.generic') {
const unresolvedOperands = [];
const allMapOperands = [];
const parseGenericOperandList = () => {
parser.parseCommaSeparatedList('none', () => {
const operand = parser.parseOperand();
unresolvedOperands.push(operand);
const mapOperands = [];
parser.parseAffineMapOfSSAIds(mapOperands, null, null, new Map());
allMapOperands.push(...mapOperands);
if (parser.parseOptionalColon()) {
parser.parseAttribute();
}
});
};
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
parseGenericOperandList();
parser.parseRParen();
}
}
if (parser.parseOptionalLess()) {
const reduction = parser.parseKeyword();
result.addAttribute('reduction', reduction);
parser.parseGreater();
}
const __kernel = parser.parseOptionalSymbolName();
if (__kernel !== null) {
result.addAttribute('kernel', __kernel);
}
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
parseGenericOperandList();
parser.parseRParen();
}
}
if (parser.parseOptionalKeyword('tile')) {
parser.parseColon();
parser.parseLSquare();
const tileValues = [];
if (!parser.parseOptionalRSquare()) {
do {
tileValues.push(parser.parseInteger());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
result.addAttribute('tile', new _.DenseI64ArrayAttr(tileValues));
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const funcType = parser.parseType();
if (funcType && funcType.inputs) {
parser.resolveOperands(unresolvedOperands, funcType.inputs, result.operands);
}
const indexType = new _.IndexType();
parser.resolveOperands(allMapOperands, allMapOperands.map(() => indexType), result.operands);
if (funcType && funcType.results) {
for (const resultType of funcType.results) {
result.addTypes([resultType]);
}
}
} else {
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, result.operands);
}
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.hlo.ToyDialect = class extends _.hlo.HLODialect {
constructor(operations) {
super(operations, 'toy');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'toy.func') {
parser.parseFunctionOp(result, false);
return true;
}
// toy.constant: {attrs} dense<...> : type
if (op === 'toy.constant') {
parser.parseOptionalAttrDict(result.attributes);
const value = parser.parseAttribute();
result.addAttribute('value', value.value === undefined ? value : value.value);
result.addTypes([value.type]);
return true;
}
// toy.mul, toy.add: %lhs, %rhs : type
if (op === 'toy.mul' || op === 'toy.add') {
result.operands = parser.parseOperandList();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const types = parser.parseOptionalColonTypeList();
if (types.length > 0) {
const [type] = types;
for (const operand of result.operands) {
operand.type = type;
}
result.addTypes(types);
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.SdfgDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'sdfg');
}
parseNumberList(parser, result, attrName) {
const operands = [];
const attrs = [];
do {
if (parser.parseOptionalKeyword('sym')) {
parser.parseLParen();
attrs.push(parser.parseAttribute());
parser.parseRParen();
} else {
const op = parser.parseOptionalOperand();
if (op) {
operands.push(op);
} else {
attrs.push(parser.parseAttribute());
}
}
} while (parser.parseOptionalComma());
if (attrName) {
result.addAttribute(attrName, attrs);
}
const indexType = new _.IndexType();
parser.resolveOperands(operands, operands.map(() => indexType), result.operands);
return operands;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'sdfg.sdfg' || op === 'sdfg.nested_sdfg' || op === 'sdir.sdfg') {
parser.parseOptionalAttrDict(result.attributes);
const inputResult = parser.parseFunctionArgumentList();
const inputs = inputResult.arguments.map((a) => a.type);
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let results = [];
if (parser.parseOptionalArrow()) {
const outputResult = parser.parseFunctionArgumentList();
results = outputResult.arguments.map((a) => a.type);
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'sdfg.tasklet' || op === 'sdir.tasklet') {
parser.parseOptionalAttrDict(result.attributes);
const __taskletSym = parser.parseOptionalSymbolName();
if (__taskletSym !== null) {
result.attributes.set('sym_name', __taskletSym);
}
const blockArgs = [];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const operand = parser.parseOperand();
let blockArgName = operand;
let type = null;
if (parser.parseOptionalKeyword('as')) {
blockArgName = parser.parseOperand();
parser.parseColon();
type = parser.parseType();
} else {
parser.parseColon();
type = parser.parseType();
}
parser.resolveOperand(operand, type, result.operands);
blockArgs.push({ value: blockArgName, type });
parser.parseOptionalComma();
}
}
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const type = parser.parseType();
result.addTypes([type]);
parser.parseOptionalComma();
}
} else {
const type = parser.parseType();
result.addTypes([type]);
}
}
const region = result.addRegion();
parser.parseRegion(region, blockArgs);
return true;
}
if (op === 'sdfg.consume') {
parser.parseOptionalAttrDict(result.attributes);
parser.parseLParen();
const stream = parser.parseOperand();
const streamType = parser.parseColonType();
parser.resolveOperand(stream, streamType, result.operands);
parser.parseRParen();
parser.parseArrow();
parser.parseLParen();
const ivs = [];
parser.parseKeyword('pe');
parser.parseColon();
const peArg = parser.parseArgument();
peArg.type = new _.IndexType();
ivs.push(peArg);
parser.parseComma();
parser.parseKeyword('elem');
parser.parseColon();
const elemArg = parser.parseArgument();
ivs.push(elemArg);
parser.parseRParen();
const region = result.addRegion();
parser.parseRegion(region, ivs);
return true;
}
if (op === 'sdfg.state' || op === 'sdir.state') {
parser.parseOptionalAttrDict(result.attributes);
const __stateSym = parser.parseOptionalSymbolName();
if (__stateSym !== null) {
result.attributes.set('sym_name', __stateSym);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'sdfg.alloc' || op === 'sdir.alloc' || op === 'sdir.alloc_transient' || op === 'sdir.alloc_stream') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedOperands = [];
if (parser.parseOptionalLParen()) {
unresolvedOperands.push(...parser.parseOperandList('none'));
parser.parseRParen();
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let allocType = null;
if (parser.parseOptionalColon()) {
allocType = parser.parseType();
result.addTypes([allocType]);
}
const indexType = new _.IndexType();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => indexType), result.operands);
return true;
}
if (op === 'sdfg.store' || op === 'sdir.store') {
parser.parseOptionalAttrDict(result.attributes);
const valueOp = parser.parseOperand();
parser.parseOptionalComma();
const arrayOp = parser.parseOperand();
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
this.parseNumberList(parser, result, 'indices');
parser.parseRSquare();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
if (parser.parseOptionalColon()) {
const valueType = parser.parseType();
parser.parseOptionalArrow();
const arrayType = parser.parseType();
parser.resolveOperand(valueOp, valueType, result.operands);
parser.resolveOperand(arrayOp, arrayType, result.operands);
} else {
parser.resolveOperand(valueOp, null, result.operands);
parser.resolveOperand(arrayOp, null, result.operands);
}
return true;
}
if (op === 'sdfg.load' || op === 'sdir.load') {
parser.parseOptionalAttrDict(result.attributes);
const arrayOp = parser.parseOperand();
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
this.parseNumberList(parser, result, 'indices');
parser.parseRSquare();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
if (parser.parseOptionalColon()) {
const arrayType = parser.parseType();
parser.parseOptionalArrow();
const resultType = parser.parseType();
parser.resolveOperand(arrayOp, arrayType, result.operands);
result.addTypes([resultType]);
} else {
parser.resolveOperand(arrayOp, null, result.operands);
}
return true;
}
if (op === 'sdfg.map' || op === 'sdir.map') {
parser.parseOptionalAttrDict(result.attributes);
const params = [];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const __param = parser.parseOptionalOperand();
if (__param) {
params.push(__param);
}
parser.parseOptionalComma();
}
}
if (parser.parseOptionalEqual()) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
this.parseNumberList(parser, result, 'lowerBounds');
parser.parseRParen();
}
}
if (parser.parseOptionalKeyword('to')) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
this.parseNumberList(parser, result, 'upperBounds');
parser.parseRParen();
}
}
if (parser.parseOptionalKeyword('step')) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
this.parseNumberList(parser, result, 'steps');
parser.parseRParen();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'sdir.consume') {
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalLParen()) {
let __sdirOp = parser.parseOptionalOperand();
while (__sdirOp) {
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(__sdirOp, type, result.operands);
if (!parser.parseOptionalComma()) {
break;
}
__sdirOp = parser.parseOptionalOperand();
}
parser.parseRParen();
}
if (parser.parseOptionalArrow()) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
parser.parseKeyword();
parser.parseColon();
parser.parseOperand();
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'sdfg.edge' || op === 'sdir.edge') {
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const __kw = parser.parseOptionalKeyword();
if (__kw) {
parser.parseColon(); // label like 'ref'
}
const __op = parser.parseOptionalOperand();
if (__op) {
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(__op, type, result.operands);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const __src = parser.parseOptionalSymbolName();
if (__src !== null) {
result.addAttribute('src', __src);
}
parser.parseOptionalArrow();
const __dst = parser.parseOptionalSymbolName();
if (__dst !== null) {
result.addAttribute('dst', __dst);
}
return true;
}
if (op === 'sdfg.sym' || op === 'sdir.sym') {
if (parser.parseOptionalLParen()) {
const expr = parser.parseString();
result.addAttribute('expr', expr);
parser.parseOptionalRParen();
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
result.addTypes([type]);
}
return true;
}
if (op === 'sdfg.copy' || op === 'sdir.copy') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedSrc = parser.parseOperandList('none');
const unresolvedDst = [];
if (parser.parseOptionalArrow()) {
unresolvedDst.push(...parser.parseOperandList('none'));
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
const allUnresolved = unresolvedSrc.concat(unresolvedDst);
parser.resolveOperands(allUnresolved, allUnresolved.map(() => type), result.operands);
} else {
const allUnresolved = unresolvedSrc.concat(unresolvedDst);
parser.resolveOperands(allUnresolved, allUnresolved.map(() => null), result.operands);
}
return true;
}
if (op === 'sdfg.libcall' || op === 'sdir.libcall') {
parser.parseOptionalAttrDict(result.attributes);
const __libname = parser.parseOptionalString();
if (__libname !== null) {
result.addAttribute('libname', __libname);
}
const unresolvedOperands = parser.parseOperandList('paren');
const types = [];
if (parser.parseOptionalColon()) {
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
types.push(parser.parseType());
parser.parseOptionalComma();
}
}
if (parser.parseOptionalArrow()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
}
}
parser.resolveOperands(unresolvedOperands, types.length > 0 ? types : unresolvedOperands.map(() => null), result.operands);
return true;
}
if (op === 'sdfg.get_access' || op === 'sdir.get_access') {
let unresolvedOperand = null;
unresolvedOperand = parser.parseOptionalOperand();
if (parser.parseOptionalColon()) {
const inputType = parser.parseType();
if (unresolvedOperand) {
parser.resolveOperand(unresolvedOperand, inputType, result.operands);
}
if (parser.parseOptionalArrow()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
}
} else if (unresolvedOperand) {
// No explicit type - try to resolve operand and infer type from its definition
parser.resolveOperand(unresolvedOperand, null, result.operands);
// If operand was resolved with a type, use it as result type
if (result.operands.length > 0 && result.operands[0].type) {
result.addTypes([result.operands[0].type]);
} else {
// Fallback: use a generic memref type
result.addTypes([new _.UnrankedMemRefType(new _.FloatType('f32'))]);
}
}
return true;
}
if (op === 'sdir.call') {
const callee = parser.parseOptionalSymbolName();
if (callee) {
result.addAttribute('callee', callee);
}
if (parser.parseOptionalLParen()) {
const unresolvedOperands = parser.parseOperandList('none');
parser.parseRParen();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
if (parser.parseOptionalColon()) {
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
parser.parseType();
parser.parseOptionalComma();
}
}
if (parser.parseOptionalArrow()) {
const resultType = parser.parseType();
result.addTypes([resultType]);
}
}
return true;
}
if (op === 'sdfg.alloc_symbol' || op === 'sdir.alloc_symbol') {
if (parser.parseOptionalLParen()) {
const sym = parser.parseString();
result.addAttribute('sym', sym);
parser.parseOptionalRParen();
}
return true;
}
if (op === 'sdfg.return') {
const unresolvedOperands = parser.parseOperandList('none');
if (unresolvedOperands.length > 0) {
const types = parser.parseOptionalColonTypeList();
parser.resolveOperands(unresolvedOperands, types, result.operands);
}
return true;
}
if (op === 'sdfg.stream_push' || op === 'sdir.stream_push') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedValue = parser.parseOperand();
parser.parseOptionalComma();
const unresolvedStream = parser.parseOperand();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let valueType = null;
let streamType = null;
if (parser.parseOptionalColon()) {
valueType = parser.parseType();
parser.parseOptionalArrow();
streamType = parser.parseType();
}
parser.resolveOperand(unresolvedValue, valueType, result.operands);
parser.resolveOperand(unresolvedStream, streamType, result.operands);
return true;
}
if (op === 'sdfg.stream_pop' || op === 'sdir.stream_pop') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedStream = parser.parseOperand();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let streamType = null;
if (parser.parseOptionalColon()) {
streamType = parser.parseType();
parser.parseOptionalArrow();
const resultType = parser.parseType();
result.addTypes([resultType]);
}
parser.resolveOperand(unresolvedStream, streamType, result.operands);
return true;
}
if (op === 'sdfg.stream_length' || op === 'sdir.stream_length') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedStream = parser.parseOperand();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let streamType = null;
if (parser.parseOptionalColon()) {
streamType = parser.parseType();
parser.parseOptionalArrow();
const resultType = parser.parseType();
result.addTypes([resultType]);
}
parser.resolveOperand(unresolvedStream, streamType, result.operands);
return true;
}
if (op === 'sdfg.view_cast' || op === 'sdir.view_cast') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedInput = parser.parseOperand();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let inputType = null;
if (parser.parseOptionalColon()) {
inputType = parser.parseType();
parser.parseOptionalArrow();
const resultType = parser.parseType();
result.addTypes([resultType]);
}
parser.resolveOperand(unresolvedInput, inputType, result.operands);
return true;
}
if (op === 'sdfg.subview' || op === 'sdir.subview') {
parser.parseOptionalAttrDict(result.attributes);
const unresolvedInput = parser.parseOperand();
parser.parseLSquare();
this.parseNumberList(parser, result, 'offsets');
parser.parseRSquare();
parser.parseLSquare();
this.parseNumberList(parser, result, 'sizes');
parser.parseRSquare();
parser.parseLSquare();
this.parseNumberList(parser, result, 'strides');
parser.parseRSquare();
parser.parseOptionalAttrDictWithKeyword(result.attributes);
let inputType = null;
if (parser.parseOptionalColon()) {
inputType = parser.parseType();
parser.parseOptionalArrow();
const resultType = parser.parseType();
result.addTypes([resultType]);
}
parser.resolveOperand(unresolvedInput, inputType, result.operands);
return true;
}
return super.parseOperation(parser, result);
}
};
_.TFLDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfl');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tfl.control_node') {
if (parser.parseOptionalLParen()) {
const unresolvedOperands = parser.parseOperandList('none');
parser.parseRParen();
// control_node operands don't have types parsed inline
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
if (parser.parseOptionalKeyword('controls')) {
const region = { blocks: [{ operations: [] }] };
const innerOp = parser.parseGenericOperation();
region.blocks[0].operations.push(innerOp);
result.regions.push(region);
// Result types are inferred from inner op results + control type
for (const opResult of innerOp.results) {
result.addTypes([opResult.type]);
}
result.addTypes([new _.Type('!tfl.control')]);
} else if (parser.parseOptionalRegion(result.addRegion())) {
const region = result.regions[result.regions.length - 1];
// Result types from yield terminator + control type
const block = region.blocks[region.blocks.length - 1];
const yieldOp = block.operations[block.operations.length - 1];
if (yieldOp && yieldOp.operands) {
for (const operand of yieldOp.operands) {
result.addTypes([operand.type]);
}
}
result.addTypes([new _.Type('!tfl.control')]);
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'tfl.add' || op === 'tfl.sub' || op === 'tfl.mul' || op === 'tfl.div' ||
op === 'tfl.floor_div' || op === 'tfl.pow' || op === 'tfl.squared_difference' ||
op === 'tfl.less' || op === 'tfl.less_equal' || op === 'tfl.greater' || op === 'tfl.greater_equal' || op === 'tfl.not_equal' ||
op === 'tfl.logical_and' || op === 'tfl.logical_or') {
// Or: (operands) <properties> : fn-type (generic form)
if (parser.parseOptionalLParen()) {
const unresolvedOperands = parser.parseOperandList('none');
parser.parseRParen();
if (parser.parseOptionalLess()) {
result.propertiesAttr = parser.parseAttribute();
parser.parseGreater();
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const fnType = parser.parseType();
if (fnType instanceof _.FunctionType) {
parser.resolveOperands(unresolvedOperands, fnType.inputs, result.operands);
result.addTypes(fnType.results);
} else {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => fnType), result.operands);
result.addTypes([fnType]);
}
} else {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
return true;
}
const unresolvedOperands = parser.parseOperandList('none');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => type), result.operands);
result.addTypes([type]);
} else {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.TFDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tf');
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'resource' || mnemonic === 'variant') {
if (parser.parseOptionalLess()) {
const subtypes = [];
do {
subtypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.Type(`!${dialect}.${mnemonic}<${subtypes.join(', ')}>`);
}
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.TFTypeDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tf_type');
this.simpleTypes = new Set([
'string', 'qint8', 'qint16', 'qint32', 'quint8', 'quint16',
'f32ref', 'f64ref', 'uint4ref', 'int4ref', 'uint8ref', 'int8ref',
'uint16ref', 'int16ref', 'uint32ref', 'int32ref', 'uint64ref', 'int64ref',
'stringref', 'boolref', 'quint8ref', 'qint8ref', 'quint16ref', 'qint16ref',
'qint32ref', 'bfloat16ref', 'complex64ref', 'complex128ref', 'halfref',
'resourceref', 'variantref',
'float8e4m3fnref', 'float8e5m2ref', 'float8e4m3fnuzref',
'float8e4m3b11fnuzref', 'float8e5m2fnuzref'
]);
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
const type = `!${dialect}.${mnemonic}`;
// Handle parametrized types like resource<>, variant<>, resource_handle<>
if (mnemonic === 'resource' || mnemonic === 'variant' || mnemonic === 'resource_handle') {
if (parser.parseOptionalLess()) {
const subtypes = [];
do {
subtypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.Type(`${type}<${subtypes.join(', ')}>`);
}
}
return new _.Type(type);
}
};
_.TransformDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'transform');
this.registerCustomDirective('PackedOrDynamicIndexList', this.parsePackedOrDynamicIndexList.bind(this));
this.registerCustomDirective('SemiFunctionType', this.parseSemiFunctionType.bind(this));
this.registerCustomDirective('SequenceOpOperands', this.parseSequenceOpOperands.bind(this));
this.registerCustomDirective('ForeachMatchSymbols', this.parseForeachMatchSymbols.bind(this));
this.registerCustomDirective('TransformMatchDims', this.parseTransformMatchDims.bind(this));
this.registerCustomDirective('ApplyRegisteredPassOptions', this.parseApplyRegisteredPassOptions.bind(this));
this.registerCustomDirective('AlternativesOpSelectedRegion', this.parseAlternativesOpSelectedRegion.bind(this));
this.registerCustomDirective('ContinuousTileSizeTypes', this.parseContinuousTileSizeTypes.bind(this));
this.registerCustomDirective('MultitileSizesTypes', this.parseMultitileSizesTypes.bind(this));
this.registerCustomDirective('InnerTileAlignmentArray', this.parseInnerTileAlignmentArray.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'transform.named_sequence') {
return this.parseNamedSequenceOp(parser, result);
}
// C++-only operation: transform.test_transform_op ["message"]
// Defined in mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
if (op === 'transform.test_transform_op') {
const __message = parser.parseOptionalString();
if (__message !== null) {
result.addAttribute('message', __message);
}
return true;
}
// LinalgTransformOps.cpp:3009 SplitOp::parse
if (op === 'transform.structured.split') {
const unresolvedTarget = parser.parseOperand();
parser.parseKeyword('after');
let unresolvedDynamicChunk = null;
unresolvedDynamicChunk = parser.parseOptionalOperand();
if (!unresolvedDynamicChunk) {
const staticChunkSizes = parser.parseInteger();
result.addAttribute('static_chunk_sizes', staticChunkSizes);
}
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const targetType = parser.parseType();
parser.resolveOperand(unresolvedTarget, targetType, result.operands);
result.addTypes([targetType]);
if (unresolvedDynamicChunk && parser.parseOptionalComma()) {
const chunkType = parser.parseType();
parser.resolveOperand(unresolvedDynamicChunk, chunkType, result.operands);
} else if (unresolvedDynamicChunk) {
// Default to index type if chunk type not specified
parser.resolveOperand(unresolvedDynamicChunk, null, result.operands);
}
return true;
}
return super.parseOperation(parser, result);
}
parseSequenceOpOperands(parser, op /*, args */) {
const unresolvedOperands = [];
let __seqOp = parser.parseOptionalOperand();
while (__seqOp) {
unresolvedOperands.push(__seqOp);
if (!parser.parseOptionalComma()) {
break;
}
__seqOp = parser.parseOptionalOperand();
}
if (parser.parseOptionalColon()) {
parser.parseOptionalLParen();
const types = parser.parseTypeList();
parser.resolveOperands(unresolvedOperands, types, op.operands);
parser.parseOptionalRParen();
} else {
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, op.operands);
}
}
}
parseForeachMatchSymbols(parser, op, matchersAttr, actionsAttr) {
const matchers = [];
const actions = [];
do {
const matcher = parser.parseOptionalSymbolName();
parser.parseArrow();
const action = parser.parseOptionalSymbolName();
matchers.push(matcher);
actions.push(action);
} while (parser.parseOptionalComma());
op.addAttribute(matchersAttr, matchers);
op.addAttribute(actionsAttr, actions);
}
parseTransformMatchDims(parser, op, dimsAttr, invertedAttr, allAttr) {
if (parser.parseOptionalKeyword('all')) {
op.addAttribute(allAttr, true);
return;
}
const isInverted = parser.parseOptionalKeyword('except');
if (isInverted) {
parser.parseLParen();
}
const dims = [];
do {
const value = parser.parseOptionalInteger();
if (value !== null) {
dims.push(value);
}
} while (parser.parseOptionalComma());
if (isInverted) {
parser.parseRParen();
op.addAttribute(invertedAttr, true);
}
op.addAttribute(dimsAttr, dims);
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
switch (mnemonic) {
case 'op': {
parser.parseLess();
const name = parser.parseAttribute();
parser.parseGreater();
return new _.Type(`!${dialect}.op<${name}>`);
}
case 'param': {
parser.parseLess();
const type = parser.parseType();
parser.parseGreater();
return new _.Type(`!${dialect}.param<${type}>`);
}
default:
return new _.Type(`!${dialect}.${mnemonic}`);
}
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseNamedSequenceOp(parser, result) {
parser.parseOptionalVisibilityKeyword(result.attributes);
parser.parseSymbolName('sym_name', result.attributes);
const argResult = parser.parseFunctionArgumentList();
const inputs = argResult.arguments.map((a) => a.type);
const results = [];
const resultAttrs = [];
if (parser.parseOptionalArrow()) {
parser.parseFunctionResultList(results, resultAttrs);
}
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType(inputs, results)));
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion());
return true;
}
parseSemiFunctionType(parser, op /* , args */) {
const hasLParen = parser.parseOptionalLParen();
const argType = parser.parseType();
if (op.operands.length > 0) {
op.operands[0].type = argType;
}
if (!hasLParen) {
return;
}
parser.parseRParen();
parser.parseArrow();
if (parser.parseOptionalLParen()) {
let idx = 0;
do {
const type = parser.parseType();
if (idx < op.types.length) {
op.types[idx] = type;
} else {
op.addTypes([type]);
}
idx++;
} while (parser.parseOptionalComma());
parser.parseRParen();
} else {
const type = parser.parseType();
if (op.types.length > 0) {
op.types[0] = type;
} else {
op.addTypes([type]);
}
}
}
parsePackedOrDynamicIndexList(parser, op, packedName, dynamicName, staticAttrName) {
const dynamicOperands = [];
const dynamicTypes = [];
const staticValues = [];
let packedOperand = null;
// Check for packed syntax: *(%operand)
if (parser.parseOptionalStar()) {
parser.parseLParen();
packedOperand = parser.parseOptionalOperand();
parser.parseRParen();
} else if (parser.parseOptionalLSquare()) {
// List syntax: [int, %operand, int, ...]
while (!parser.parseOptionalRSquare()) {
const __dynOp = parser.parseOptionalOperand();
if (__dynOp) {
dynamicOperands.push(__dynOp);
staticValues.push(-9223372036854775808n); // ShapedType::kDynamic
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
dynamicTypes.push(type);
} else {
const __intVal = parser.parseOptionalInteger('int64');
if (__intVal === null) {
break;
}
staticValues.push(__intVal);
}
parser.parseOptionalComma();
}
}
if (packedOperand && packedName) {
parser.resolveOperand(packedOperand, null, op.operands);
}
if (dynamicName) {
for (let i = 0; i < dynamicOperands.length; i++) {
const type = i < dynamicTypes.length ? dynamicTypes[i] : null;
parser.resolveOperand(dynamicOperands[i], type, op.operands);
}
}
if (staticAttrName && staticValues.length > 0) {
op.addAttribute(staticAttrName, staticValues);
}
}
parseContinuousTileSizeTypes(parser, result, targetTypes, tileSizesTypes, chunkSizesTypes) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
if (funcType.inputs.length !== 1 || funcType.results.length !== 1) {
throw new mlir.Error(`Expected a trailing functional type with one argument and one result.`);
}
targetTypes.push(funcType.inputs[0]);
tileSizesTypes.push(funcType.results[0]);
chunkSizesTypes.push(funcType.results[0]);
}
}
parseMultitileSizesTypes(parser, result, targetTypes, lowSizeTypes, highSizeTypes, splitPointTypes) {
const funcType = parser.parseType();
if (funcType instanceof _.FunctionType) {
if (funcType.inputs.length !== 1 || funcType.results.length !== 1) {
throw new mlir.Error(`Expected a trailing functional type with one argument and one result.`);
}
targetTypes.push(funcType.inputs[0]);
lowSizeTypes.push(funcType.results[0]);
highSizeTypes.push(funcType.results[0]);
splitPointTypes.push(funcType.results[0]);
}
}
parseApplyRegisteredPassOptions(parser, op, optionsAttr, dynamicOptions) {
if (!parser.parseOptionalLBrace()) {
return;
}
const options = {};
let dynamicIdx = 0;
while (!parser.parseOptionalRBrace()) {
let key = parser.parseOptionalString();
if (key === null) {
key = parser.parseOptionalKeyword();
}
parser.parseEqual();
const __regOp = parser.parseOptionalOperand();
if (__regOp) {
dynamicOptions.push(__regOp);
options[key] = `#transform.param_operand<${dynamicIdx++}>`;
} else if (parser.parseOptionalLSquare()) {
const arr = [];
while (!parser.parseOptionalRSquare()) {
const __arrOp = parser.parseOptionalOperand();
if (__arrOp) {
dynamicOptions.push(__arrOp);
arr.push(`#transform.param_operand<${dynamicIdx++}>`);
} else {
const val = parser.parseAttribute();
arr.push(val);
}
parser.parseOptionalComma();
}
options[key] = arr;
} else {
const value = parser.parseAttribute();
options[key] = value;
}
parser.parseOptionalComma();
}
op.addAttribute(optionsAttr, options);
}
parseAlternativesOpSelectedRegion(parser, op, selectedRegionAttr, selectedRegionParam) {
const __altInt = parser.parseOptionalInteger();
if (__altInt === null) {
const __altOp = parser.parseOptionalOperand();
if (__altOp) {
selectedRegionParam.push(__altOp);
}
} else {
op.addAttribute(selectedRegionAttr, __altInt);
}
}
parseInnerTileAlignmentArray(parser, op, attrName) {
const alignments = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
const alignment = parser.parseOptionalKeyword();
if (alignment === null) {
break;
}
alignments.push(alignment);
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (attrName) {
op.addAttribute(attrName, alignments);
}
}
};
_.test = {};
_.test.TestDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'test');
this.blobManager = new Map();
this.registerCustomDirective('CustomOptionalOperand', this.parseCustomOptionalOperand.bind(this));
this.registerCustomDirective('CustomDirectiveOperands', this.parseCustomDirectiveOperands.bind(this));
this.registerCustomDirective('CustomDirectiveOperandsAndTypes', this.parseCustomDirectiveOperandsAndTypes.bind(this));
this.registerCustomDirective('CustomDirectiveResults', this.parseCustomDirectiveResults.bind(this));
this.registerCustomDirective('CustomDirectiveWithTypeRefs', this.parseCustomDirectiveWithTypeRefs.bind(this));
this.registerCustomDirective('CustomDirectiveRegions', this.parseCustomDirectiveRegions.bind(this));
this.registerCustomDirective('CustomDirectiveSuccessors', this.parseCustomDirectiveSuccessors.bind(this));
this.registerCustomDirective('CustomDirectiveAttrDict', this.parseCustomDirectiveAttrDict.bind(this));
this.registerCustomDirective('CustomDirectiveAttributes', this.parseCustomDirectiveAttributes.bind(this));
this.registerCustomDirective('CustomDirectiveSpacing', this.parseCustomDirectiveSpacing.bind(this));
this.registerCustomDirective('CustomDirectiveOptionalOperandRef', this.parseCustomDirectiveOptionalOperandRef.bind(this));
this.registerCustomDirective('UsingPropertyInCustom', this.parseUsingPropertyInCustom.bind(this));
this.registerCustomDirective('IntProperty', this.parseIntProperty.bind(this));
this.registerCustomDirective('SumProperty', this.parseSumProperty.bind(this));
this.registerCustomDirective('SwitchCases', this.parseSwitchCases.bind(this));
this.registerCustomDirective('DimensionList', this.parseDimensionList.bind(this));
this.registerCustomDirective('OptionalCustomParser', this.parseOptionalCustomParser.bind(this));
this.registerCustomDirective('OptionalLoc', this.parseOptionalLoc.bind(this));
this.registerCustomDirective('DummyRegionRef', this.parseDummyRegionRef.bind(this));
this.registerCustomDirective('DummySuccessorRef', this.parseDummySuccessorRef.bind(this));
this.registerCustomType('CompoundNestedOuterType', this.parseCompoundNestedOuterType.bind(this));
this.registerCustomType('CompoundNestedInnerType', this.parseCompoundNestedInnerType.bind(this));
this.registerCustomType('CompoundTypeA', this.parseCompoundTypeA.bind(this));
this.registerCustomAttribute('TestBitEnumAttr', this.parseEnumFlagsAngleBracketComma.bind(this));
this.registerCustomAttribute('TestBitEnumVerticalBarAttr', this.parseEnumFlagsAngleBracketPipe.bind(this));
this.registerCustomAttribute('TestEnumAttr', this.parseTestEnumAttr.bind(this));
this.registerCustomAttribute('TestEnumProp', this.parseTestEnumAttr.bind(this));
this.registerCustomAttribute('TestEnumPropAttrForm', this.parseTestEnumPropAttrForm.bind(this));
this.registerCustomAttribute('TestBitEnumProp', (parser, type) => this.parseEnumFlags(parser, type, ',', true));
this.registerCustomAttribute('TestBitEnumPropNamed', this.parseTestBitEnumPropNamed.bind(this));
this.registerCustomAttribute('TestArrayOfUglyAttrs', this.parseTestArrayOfUglyAttrs.bind(this));
this.registerCustomAttribute('TestArrayOfInts', this.parseTestArrayOfInts.bind(this));
this.registerCustomAttribute('TestArrayOfEnums', this.parseTestArrayOfEnums.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'test.conversion_func_op') {
parser.parseFunctionOp(result, false);
return true;
}
if (op === 'test.region_if') {
const unresolvedOperands = [];
let __regionOp = parser.parseOptionalOperand();
while (__regionOp) {
unresolvedOperands.push(__regionOp);
if (!parser.parseOptionalComma()) {
break;
}
__regionOp = parser.parseOptionalOperand();
}
parser.parseColon();
const inputTypes = parser.parseTypeList();
parser.resolveOperands(unresolvedOperands, inputTypes, result.operands);
const outputTypes = parser.parseArrowTypeList();
for (const t of outputTypes) {
result.addTypes([t]);
}
parser.parseKeyword('then');
const thenRegion = {};
parser.parseRegion(thenRegion);
result.regions.push(thenRegion);
parser.parseKeyword('else');
const elseRegion = {};
parser.parseRegion(elseRegion);
result.regions.push(elseRegion);
parser.parseKeyword('join');
const joinRegion = {};
parser.parseRegion(joinRegion);
result.regions.push(joinRegion);
return true;
}
if (op === 'test.affine_scope' || op === 'test.single_no_terminator_custom_asm_op') {
const region = result.addRegion();
parser.parseRegion(region);
return true;
}
if (op === 'test.with_nice_properties') {
result.compatibility = true;
let label = null;
label = parser.parseOptionalString();
if (label === null) {
label = parser.parseKeyword();
}
parser.parseKeyword('is');
const negative = parser.parseOptionalMinus();
const value = parser.parseInteger();
result.addAttribute('prop', { label, value: negative ? -value : value });
parser.parseOptionalAttrDict(result.attributes);
return true;
}
// Test operation with default-valued properties and UnitProp
if (op === 'test.with_default_valued_properties') {
result.compatibility = true;
if (parser.parseOptionalKeyword('na')) {
// All defaults
} else {
const a = parser.parseInteger();
result.addAttribute('a', a);
const __bVal = parser.parseOptionalString();
if (__bVal !== null) {
result.addAttribute('b', __bVal);
}
const c = parser.parseOptionalInteger();
if (c !== null) {
result.addAttribute('c', c);
}
if (parser.parseOptionalKeyword('unit')) {
result.addAttribute('unit', true);
} else if (parser.parseOptionalKeyword('unit_absent')) {
result.addAttribute('unit', false);
}
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
// Test operation with optional properties using some<...> syntax
if (op === 'test.with_optional_properties') {
result.compatibility = true;
const parseOptionalValue = () => {
if (parser.parseOptionalKeyword('some')) {
parser.parseLess();
let value = null;
if (parser.parseOptionalKeyword('none')) {
value = null;
} else if (parser.parseOptionalKeyword('unit')) {
value = true;
} else {
const __strVal = parser.parseOptionalString();
if (__strVal === null) {
const neg = parser.parseOptionalMinus();
value = parser.parseInteger();
if (neg) {
value = -value;
}
} else {
value = __strVal;
}
}
parser.parseGreater();
return { some: value };
}
const __retval2 = parser.parseOptionalString();
if (__retval2 !== null) {
return __retval2;
}
const __neg = parser.parseOptionalMinus();
const __val = parser.parseInteger();
return __neg ? -__val : __val;
};
const knownAttrs = ['anAttr', 'simple', 'simplei8', 'simpleui8', 'nonTrivialStorage', 'hasDefault', 'nested', 'longSyntax', 'hasUnit', 'maybeUnit'];
for (;;) {
const name = parser.parseOptionalKeyword(knownAttrs);
if (!name) {
break;
}
if (name === 'hasUnit') {
result.addAttribute(name, true);
} else if (parser.parseOptionalEqual()) {
result.addAttribute(name, parseOptionalValue());
} else {
break;
}
}
parser.parseOptionalAttrDict(result.attributes);
return true;
}
if (op === 'test.wrapping_region') {
parser.parseKeyword('wraps');
const region = result.addRegion();
const block = { operations: [] };
region.blocks = [block];
const wrappedOp = parser.parseGenericOperation();
block.operations.push(wrappedOp);
if (wrappedOp.results) {
for (const opResult of wrappedOp.results) {
result.addTypes([opResult.type]);
}
}
return true;
}
if (op === 'test.pretty_printed_region') {
result.operands = parser.parseOperandList();
if (parser.parseOptionalKeyword('start')) {
const innerOpName = parser.parseCustomOperationName();
result.addAttribute('inner_op', innerOpName);
parser.parseKeyword('end');
parser.parseColon();
const fnType = parser.parseType();
if (fnType && fnType.inputs) {
parser.resolveOperands(result.operands, fnType.inputs);
}
if (fnType && fnType.results) {
for (let i = 0; i < fnType.results.length; i++) {
result.addTypes([fnType.results[i]]);
}
}
parser.parseOptionalLocationSpecifier();
} else {
parser.parseLParen();
const region = result.addRegion();
parser.parseRegion(region);
parser.parseRParen();
parser.parseColon();
const fnType = parser.parseType();
if (fnType && fnType.inputs) {
parser.resolveOperands(result.operands, fnType.inputs);
}
if (fnType && fnType.results) {
for (let i = 0; i < fnType.results.length; i++) {
result.addTypes([fnType.results[i]]);
}
}
}
return true;
}
if (op === 'test.isolated_region') {
const operand = parser.parseOperand();
const indexType = new _.IndexType();
parser.resolveOperand(operand, indexType, result.operands);
const region = result.addRegion();
parser.parseRegion(region, [{ value: operand.toString(), type: new _.IndexType() }]);
return true;
}
if (op === 'test.string_attr_pretty_name') {
const numResults = parser.getNumResults();
for (let i = 0; i < numResults; i++) {
result.addTypes([new _.IntegerType('i32')]);
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
if (!result.attributes.has('names') && numResults > 0) {
const names = [];
for (let i = 0; i < numResults; i++) {
const resultName = parser.getResultName(i);
// Only use names that don't start with a digit (ref impl line 572)
if (resultName && resultName.length > 0 && !/^\d/.test(resultName)) {
names.push(new _.StringAttr(resultName));
} else {
names.push(new _.StringAttr(''));
}
}
result.attributes.set('names', new _.ArrayAttr(names));
}
return true;
}
if (op === 'test.with_bounds_region') {
parser.parseOptionalAttrDict(result.attributes);
const argName = parser.parseOperand();
parser.parseColon();
const argType = parser.parseType();
const region = result.addRegion();
const arg = { value: argName, type: argType };
parser.parseRegion(region, [arg]);
return true;
}
return super.parseOperation(parser, result);
}
parseTestEnumAttr(parser, type) {
const keyword = parser.parseOptionalKeyword(type.values);
if (keyword) {
return new _.TypedAttr(keyword, null);
}
return null;
}
parseTestEnumPropAttrForm(parser) {
return parser.parseOptionalAttribute();
}
parseTestBitEnumPropNamed(parser) {
if (parser.parseOptionalKeyword('bit_enum')) {
if (parser.parseOptionalLess()) {
const flags = [];
do {
const value = parser.parseKeyword();
flags.push(value);
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.TypedAttr(`bit_enum<${flags.join(', ')}>`, null);
}
}
return null;
}
// Parse CompoundNestedOuterType: assemblyFormat = "`<` `i` $inner `>`"
// Full form: !test.cmpnd_nested_outer<i !test.cmpnd_inner<...>>
// Elided form: <i <...>>
parseCompoundNestedOuterType(parser) {
parser.parseLess();
parser.parseKeyword('i');
// Parse $inner - could be full (!test.cmpnd_inner<...>) or elided (<...>)
const inner = parser.parseOptionalType() || this.parseCompoundNestedInnerType(parser);
parser.parseGreater();
return new _.Type(`!test.cmpnd_nested_outer<i ${inner}>`);
}
// Parse CompoundNestedInnerType: assemblyFormat = "`<` $some_int $cmpdA `>`"
// Full form: !test.cmpnd_inner<42 !test.cmpnd_a<...>>
// Elided form: <42 <...>>
parseCompoundNestedInnerType(parser) {
parser.parseLess();
const someInt = parser.parseInteger();
// Parse $cmpdA - could be full (!test.cmpnd_a<...>) or elided (<...>)
const cmpdA = parser.parseOptionalType() || this.parseCompoundTypeA(parser);
parser.parseGreater();
return new _.Type(`!test.cmpnd_inner<${someInt} ${cmpdA}>`);
}
// Parse CompoundTypeA: hasCustomAssemblyFormat = 1
// Example: <1, !test.smpla, [5, 6]>
parseCompoundTypeA(parser) {
parser.parseLess();
const width = parser.parseInteger();
parser.parseComma();
const oneType = parser.parseType();
parser.parseComma();
parser.parseLSquare();
const arrayOfInts = [];
if (!parser.parseOptionalRSquare()) {
do {
arrayOfInts.push(parser.parseInteger());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parser.parseGreater();
return new _.Type(`!test.cmpnd_a<${width}, ${oneType}, [${arrayOfInts.join(', ')}]>`);
}
parseOptionalLoc(parser, op, attrName = 'loc') {
const loc = parser.parseOptionalLocationSpecifier();
if (loc) {
op.addAttribute(attrName, loc);
} else {
op.addAttribute(attrName, parser.getCurrentLocation());
}
}
parseDummyRegionRef() {
}
parseDummySuccessorRef() {
}
parseTestAttrUgly(parser) {
parser.parseKeyword('begin');
const attr = parser.parseAttribute();
parser.parseKeyword('end');
return attr;
}
// Parse TestArrayOfUglyAttrs: assemblyFormat = "`[` (`]`) : ($value^ ` ` `]`)?"
parseTestArrayOfUglyAttrs(parser) {
const elements = [];
parser.parseCommaSeparatedList('square', () => {
elements.push(this.parseTestAttrUgly(parser));
});
return new _.ArrayAttr(elements);
}
// Parse TestArrayOfInts: assemblyFormat = "`[` (`]`) : ($value^ `]`)?"
parseTestArrayOfInts(parser) {
const elements = [];
parser.parseCommaSeparatedList('square', () => {
const isNegative = parser.parseOptionalMinus();
const value = parser.parseInteger();
elements.push(new _.IntegerAttr(null, isNegative ? -value : value));
});
return new _.ArrayAttr(elements);
}
// Parse TestArrayOfEnums: elements are SimpleEnumAttr values (a, b, etc.)
parseTestArrayOfEnums(parser) {
const elements = [];
parser.parseCommaSeparatedList('square', () => {
const __enumStr = parser.parseOptionalString();
if (__enumStr === null) {
const value = parser.parseKeyword();
elements.push(new _.TypedAttr(value, null));
} else {
elements.push(new _.TypedAttr(__enumStr, null));
}
});
return new _.ArrayAttr(elements);
}
parseOptionalCustomParser(parser, op, attrName = 'attr') {
if (!parser.parseOptionalKeyword('foo')) {
return null; // Optional group not taken
}
const attr = parser.parseAttribute();
op.addAttribute(attrName, attr.value);
return true;
}
parseDimensionList(parser, op, attrName = 'dimension_list') {
if (parser.parseOptionalLSquare()) {
parser.parseOptionalRSquare();
op.addAttribute(attrName, []);
return;
}
const dimInfo = parser.parseDimensionList(true, false);
op.addAttribute(attrName, dimInfo.dimensions);
}
parseCustomOptionalOperand(parser, op, optOperand) {
if (parser.parseOptionalLParen()) {
optOperand.push(parser.parseOperand());
parser.parseRParen();
}
}
// Custom directive: operand [, optOperand] -> (varOperands)
parseCustomDirectiveOperands(parser, op, operand, optOperand, varOperands) {
operand.push(parser.parseOperand());
if (parser.parseOptionalComma()) {
optOperand.push(parser.parseOperand());
}
parser.parseArrow();
parser.parseLParen();
let __varOp = parser.parseOptionalOperand();
while (__varOp) {
varOperands.push(__varOp);
if (!parser.parseOptionalComma()) {
break;
}
__varOp = parser.parseOptionalOperand();
}
parser.parseRParen();
}
parseCustomDirectiveOperandsAndTypes(parser, op, operand, optOperand, varOperands, operandTypes, optOperandTypes, varOperandTypes) {
this.parseCustomDirectiveOperands(parser, op, operand, optOperand, varOperands);
this.parseCustomDirectiveResults(parser, op, operandTypes, optOperandTypes, varOperandTypes);
}
parseCustomDirectiveResults(parser, op, operandType, optOperandType, varOperandTypes) {
parser.parseColon();
operandType.push(parser.parseType());
if (parser.parseOptionalComma()) {
optOperandType.push(parser.parseType());
}
parser.parseArrow();
parser.parseLParen();
const types = parser.parseTypeList();
for (const t of types) {
varOperandTypes.push(t);
}
parser.parseRParen();
}
parseCustomDirectiveWithTypeRefs(parser, op /*, operandType, optOperandType, varOperandTypes */) {
parser.parseKeyword('type_refs_capture');
const operandType2 = [];
const optOperandType2 = [];
const varOperandTypes2 = [];
this.parseCustomDirectiveResults(parser, op, operandType2, optOperandType2, varOperandTypes2);
}
parseCustomDirectiveRegions(parser, op /*, regionName, varRegionsName */) {
const region = op.addRegion();
parser.parseRegion(region);
if (parser.parseOptionalComma()) {
const varRegion = op.addRegion();
parser.parseRegion(varRegion);
}
}
parseCustomDirectiveSuccessors(parser, op /*, successorName, varSuccessorsName */) {
if (!op.successors) {
op.successors = [];
}
const successor = parser.parseSuccessor();
op.successors.push(successor);
if (parser.parseOptionalComma()) {
const varSuccessor = parser.parseSuccessor();
op.successors.push(varSuccessor);
op.successors.push(varSuccessor);
}
}
parseCustomDirectiveAttrDict(parser, op) {
parser.parseOptionalAttrDict(op.attributes);
}
parseCustomDirectiveAttributes(parser, op, attrName, optAttrName) {
const attr = parser.parseAttribute();
op.addAttribute(attrName, attr);
if (parser.parseOptionalComma()) {
const optAttr = parser.parseAttribute();
op.addAttribute(optAttrName, optAttr);
}
}
parseCustomDirectiveSpacing(parser, op, attrName) {
if (attrName) {
const name = attrName.name || attrName;
const attr = parser.parseAttribute();
op.addAttribute(name, attr);
}
}
parseCustomDirectiveOptionalOperandRef(parser) {
parser.parseInteger();
}
parseSwitchCases(parser, op, casesAttrName /*, caseRegionsName */) {
const caseValues = [];
while (parser.parseOptionalKeyword('case')) {
const value = parser.parseInteger();
caseValues.push(value);
const region = op.addRegion();
parser.parseRegion(region);
}
if (casesAttrName) {
op.addAttribute(casesAttrName, caseValues);
}
}
parseUsingPropertyInCustom(parser, op, propArg) {
const values = [];
parser.parseLSquare();
if (!parser.parseOptionalRSquare()) {
do {
values.push(parser.parseInteger());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
if (propArg) {
const propName = typeof propArg === 'string' ? propArg : propArg.name;
op.addAttribute(propName, `array<i64: ${values.join(', ')}>`);
}
}
parseIntProperty(parser, op, propArg) {
const value = parser.parseInteger();
if (propArg) {
const propName = typeof propArg === 'string' ? propArg : propArg.name;
op.addAttribute(propName, value);
}
}
parseSumProperty(parser, op, propArg) {
const second = parser.parseInteger();
parser.parseEqual();
parser.parseInteger(); // sum value (validation skipped)
if (propArg) {
const propName = typeof propArg === 'string' ? propArg : propArg.name;
op.addAttribute(propName, second);
}
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'test.format_infer_type' || op.name.getStringRef() === 'test.format_infer_type2') {
op.addTypes([new _.IntegerType('i16')]);
return;
}
if (op.name.getStringRef() === 'test.format_types_match_attr') {
const valueAttr = op.attributes.get('value');
if (valueAttr && valueAttr.type) {
op.addTypes([valueAttr.type]);
return;
}
}
if (op.name.getStringRef() === 'test.format_infer_type_all_operands_and_types' ||
op.name.getStringRef() === 'test.format_infer_type_all_types_one_operand' ||
op.name.getStringRef() === 'test.format_infer_type_all_types_two_operands' ||
op.name.getStringRef() === 'test.format_infer_type_all_types' ||
op.name.getStringRef() === 'test.format_infer_type_variadic_operands') {
for (const operand of op.operands) {
if (operand.type) {
op.addTypes([operand.type]);
}
}
return;
}
if (op.name.getStringRef() === 'test.format_infer_type_regions') {
const region = op.regions[0];
const block = region?.blocks?.[0];
if (block?.arguments) {
for (const arg of block.arguments) {
if (arg.type) {
op.addTypes([arg.type]);
}
}
}
return;
}
if (op.name.getStringRef() === 'test.with_properties_and_inferred_type') {
const lhsAttr = op.attributes.get('lhs');
const rhsAttr = op.getAttr('rhs');
const lhs = lhsAttr && lhsAttr.value !== undefined ? lhsAttr.value : (lhsAttr || 0);
const rhs = rhsAttr !== undefined && rhsAttr !== null ? rhsAttr : 0;
const width = (typeof lhs === 'number' ? lhs : 0) + (typeof rhs === 'number' ? rhs : 0);
op.addTypes([new _.IntegerType(`i${width}`)]);
return;
}
super.inferResultTypes(op, vars);
}
declareResource(key) {
if (!this.blobManager.has(key)) {
this.blobManager.set(key, new _.DenseResourceElementsHandle(key));
}
return this.blobManager.get(key);
}
getResourceKey(handle) {
return handle.key;
}
parseResource(entry) {
const blob = entry.parseAsBlob();
this.blobManager.get(entry.key).blob = blob;
}
};
_.triton = {};
_.triton.nvws = {};
_.triton.gluon = {};
_.triton.nvidia_gpu = {};
_.triton.amdgpu = {};
_.triton.proton = {};
_.triton.xla = {};
_.triton.TensorDescType = class extends _.Type {
constructor(shape, elementType, sharedLayout) {
super(null);
this.shape = shape;
this.elementType = elementType;
this.sharedLayout = sharedLayout || null;
}
static parse(parser) {
parser.parseLess();
const blockType = parser.parseOptionalType();
if (blockType instanceof _.RankedTensorType) {
parser.parseGreater();
return new _.triton.TensorDescType(blockType.shape, blockType.elementType, blockType.encoding);
}
const dimInfo = parser.parseDimensionList(false, true);
const elementType = parser.parseType();
let sharedLayout = null;
if (parser.parseOptionalComma()) {
sharedLayout = parser.parseAttribute();
}
parser.parseGreater();
return new _.triton.TensorDescType(dimInfo.dimensions, elementType, sharedLayout);
}
toString() {
const shapeStr = this.shape.map((d) => `${d}x`).join('');
const parts = [`${shapeStr}${this.elementType}`];
if (this.sharedLayout) {
parts.push(`${this.sharedLayout}`);
}
return `!tt.tensordesc<${parts.join(', ')}>`;
}
};
_.triton.PointerType = class extends _.Type {
constructor(pointeeType, addressSpace) {
super(null);
this.pointeeType = pointeeType;
this.addressSpace = addressSpace || 1;
}
static parse(parser) {
parser.parseLess();
const pointeeType = parser.parseType();
let addressSpace = 1;
if (parser.parseOptionalComma()) {
addressSpace = parser.parseInteger();
}
parser.parseGreater();
return new _.triton.PointerType(pointeeType, addressSpace);
}
toString() {
const pointeeStr = this.pointeeType?.toString ? this.pointeeType.toString() : this.pointeeType;
if (this.addressSpace && this.addressSpace !== 1) {
return `!tt.ptr<${pointeeStr}, ${this.addressSpace}>`;
}
return `!tt.ptr<${pointeeStr}>`;
}
};
_.triton.TritonDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tt');
this.registerCustomType('TT_Ptr', (parser) => _.triton.PointerType.parse(parser));
this.registerCustomType('TT_TensorDescType', this.parseTensorDescType.bind(this));
this.registerCustomType('TT_TensorPtr', (parser) => _.triton.PointerType.parse(parser));
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'ptr') {
return _.triton.PointerType.parse(parser);
}
if (mnemonic === 'tensordesc') {
return _.triton.TensorDescType.parse(parser);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tt.func') {
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
parseTensorDescType(parser) {
if (parser.parseOptionalLess()) {
const blockType = parser.parseOptionalType();
if (blockType instanceof _.RankedTensorType) {
parser.parseGreater();
return new _.triton.TensorDescType(blockType.shape, blockType.elementType, blockType.encoding);
}
const dimInfo = parser.parseDimensionList(false, true);
const elementType = parser.parseType();
let sharedLayout = null;
if (parser.parseOptionalComma()) {
sharedLayout = parser.parseAttribute();
}
parser.parseGreater();
return new _.triton.TensorDescType(dimInfo.dimensions, elementType, sharedLayout);
}
return null;
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'tt.load' && op.operands.length > 0) {
const ptrType = op.operands[0].type;
if (ptrType instanceof _.triton.PointerType) {
op.addTypes([ptrType.pointeeType]);
} else if (ptrType instanceof _.RankedTensorType) {
const elementType = ptrType.elementType;
if (elementType instanceof _.triton.PointerType) {
const pointeeType = elementType.pointeeType;
op.addTypes([new _.RankedTensorType(ptrType.shape, pointeeType, ptrType.encoding)]);
return;
}
}
}
if (op.name.getStringRef() === 'tt.unsplat') {
const srcType = op.operands[0]?.type;
if (srcType instanceof _.RankedTensorType) {
op.addTypes([srcType.elementType]);
return;
}
}
super.inferResultTypes(op, vars);
}
};
_.triton.gpu = {};
_.triton.gpu.MemDescType = class extends _.Type {
constructor(shape, elementType, encoding, memorySpace, mutableMemory, allocShape) {
super(null);
this.shape = shape;
this.elementType = elementType;
this.encoding = encoding;
this.memorySpace = memorySpace;
this.mutableMemory = mutableMemory || false;
this.allocShape = allocShape || null;
}
static parse(parser) {
parser.parseLess();
const dimInfo = parser.parseDimensionList();
const elementType = parser.parseType();
let encoding = null;
let memorySpace = null;
let mutableMemory = false;
let allocShape = null;
if (parser.parseOptionalComma()) {
encoding = parser.parseAttribute();
if (parser.parseOptionalComma()) {
memorySpace = parser.parseAttribute();
if (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('mutable')) {
mutableMemory = true;
if (parser.parseOptionalComma()) {
allocShape = parser.parseDimensionList(false, false).dimensions;
}
} else {
allocShape = parser.parseDimensionList(false, false).dimensions;
}
}
}
}
parser.parseGreater();
return new _.triton.gpu.MemDescType(dimInfo.dimensions, elementType, encoding, memorySpace, mutableMemory, allocShape);
}
toString() {
const shapeStr = this.shape.map((d) => `${d}x`).join('');
const parts = [`${shapeStr}${this.elementType}`];
if (this.encoding) {
parts.push(`${this.encoding}`);
}
if (this.memorySpace) {
parts.push(`${this.memorySpace}`);
}
if (this.mutableMemory) {
parts.push('mutable');
}
if (this.allocShape && this.allocShape.join('x') !== this.shape.join('x')) {
parts.push(this.allocShape.join('x'));
}
return `!ttg.memdesc<${parts.join(', ')}>`;
}
};
_.triton.gpu.TritonGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'ttg');
this.registerCustomType('TTG_MemDescType', (parser) => _.triton.gpu.MemDescType.parse(parser));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'ttg.warp_specialize') {
const unresolvedOperands = [];
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
unresolvedOperands.push(parser.parseOperand());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseKeyword('default');
const defaultRegion = {};
parser.parseRegion(defaultRegion);
result.regions.push(defaultRegion);
const partitionNumWarps = [];
let partitionIndex = 0;
while (parser.parseOptionalKeyword(`partition${partitionIndex}`)) {
const argResult = parser.parseFunctionArgumentList();
parser.parseKeyword('num_warps');
parser.parseLParen();
const numWarps = parser.parseInteger();
partitionNumWarps.push(parseInt(numWarps, 10));
parser.parseRParen();
const partitionRegion = {};
partitionRegion.arguments = argResult.arguments;
parser.parseRegion(partitionRegion);
if (!result.regions[1]) {
result.regions[1] = { blocks: [{ operations: [] }] };
}
partitionIndex++;
}
parser.parseColon();
const fnType = parser.parseType();
if (fnType instanceof _.FunctionType) {
result.addAttribute('function_type', new _.TypeAttrOf(fnType));
parser.resolveOperands(unresolvedOperands, fnType.inputs, result.operands);
result.addTypes(fnType.results);
} else {
result.addAttribute('function_type', new _.TypeAttrOf(new _.FunctionType([], [fnType])));
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
result.addTypes([fnType]);
}
if (partitionNumWarps.length > 0) {
result.addAttribute('partitionNumWarps', { type: 'array', element_type: 'i32', value: partitionNumWarps });
}
return true;
}
if (op === 'ttg.barrier') {
const flags = [];
flags.push(parser.parseKeyword());
while (parser.parseOptionalVerticalBar()) {
flags.push(parser.parseKeyword());
}
result.addAttribute('addrSpace', new _.TypedAttr(flags.join('|')));
return true;
}
return super.parseOperation(parser, result);
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'memdesc') {
return _.triton.gpu.MemDescType.parse(parser);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
};
_.triton.gluon.GluonDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'gluon');
}
};
_.triton.nvidia_gpu.TensorDescIm2ColType = class extends _.Type {
constructor(shape, elementType, sharedLayout) {
super(null);
this.shape = shape;
this.elementType = elementType;
this.sharedLayout = sharedLayout || null;
}
static parse(parser) {
parser.parseLess();
const blockType = parser.parseOptionalType();
if (blockType instanceof _.RankedTensorType) {
parser.parseGreater();
return new _.triton.nvidia_gpu.TensorDescIm2ColType(blockType.shape, blockType.elementType, blockType.encoding);
}
const dimInfo = parser.parseDimensionList(false, true);
const elementType = parser.parseType();
let sharedLayout = null;
if (parser.parseOptionalComma()) {
sharedLayout = parser.parseAttribute();
}
parser.parseGreater();
return new _.triton.nvidia_gpu.TensorDescIm2ColType(dimInfo.dimensions, elementType, sharedLayout);
}
toString() {
const shapeStr = this.shape.map((d) => `${d}x`).join('');
const parts = [`${shapeStr}${this.elementType}`];
if (this.sharedLayout) {
parts.push(`${this.sharedLayout}`);
}
return `!ttng.tensordesc_im2col<${parts.join(', ')}>`;
}
};
_.triton.nvidia_gpu.TritonNvidiaGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'ttng');
this.registerCustomDirective('Token', this.parseToken.bind(this));
this.registerCustomDirective('BarriersAndPreds', this.parseBarriersAndPreds.bind(this));
}
parseToken(parser, op, depOperands, tokenTypeArr) {
if (!parser.parseOptionalLSquare()) {
return;
}
if (Array.isArray(tokenTypeArr)) {
tokenTypeArr.push(new _.Type('!ttng.async.token'));
}
if (parser.parseOptionalRSquare()) {
return;
}
const __depOp = parser.parseOptionalOperand();
if (__depOp) {
if (!Array.isArray(depOperands)) {
parser.emitError('Expected depOperands to be an array');
}
depOperands.push(__depOp);
}
parser.parseRSquare();
}
parseBarriersAndPreds(parser, op, barrierOperands, predOperands) {
while (parser.parseOptionalComma()) {
const __barrier = parser.parseOptionalOperand();
if (__barrier) {
if (!Array.isArray(barrierOperands)) {
parser.emitError('Expected barrierOperands to be an array');
}
barrierOperands.push(__barrier);
if (parser.parseOptionalLSquare()) {
const __pred = parser.parseOptionalOperand();
if (__pred) {
if (!Array.isArray(predOperands)) {
parser.emitError('Expected predOperands to be an array');
}
predOperands.push(__pred);
}
parser.parseRSquare();
}
}
}
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'tensordesc_im2col') {
return _.triton.nvidia_gpu.TensorDescIm2ColType.parse(parser);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
};
_.triton.amdgpu.TritonAMDGPUDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'amdg');
this.registerCustomType('TT_Ptr', 'tt.ptr');
this.registerCustomType('TT_TensorPtr', 'tt.ptr');
this.registerCustomType('TTG_MemDescType', 'ttg.memdesc');
}
};
_.triton.proton.ProtonDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'proton');
}
};
_.triton.nvws.ArefType = class extends _.Type {
constructor(baseTypes) {
super(null);
this.baseTypes = baseTypes;
}
static parse(parser) {
parser.parseLess();
parser.parseLSquare();
const types = [];
if (!parser.parseOptionalRSquare()) {
do {
types.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parser.parseGreater();
return new _.triton.nvws.ArefType(types);
}
toString() {
return `!nvws.aref<[${this.baseTypes.join(', ')}]>`;
}
};
_.triton.nvws.NVWSDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'nvws');
this.registerCustomType('NVWS_ArefType', (parser) => _.triton.nvws.ArefType.parse(parser));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'nvws.warp_group') {
parser.parseOptionalAttrDictWithKeyword(result.attributes);
const numWarps = [];
let partitionIndex = 0;
while (parser.parseOptionalKeyword(`partition${partitionIndex}`)) {
parser.parseKeyword('num_warps');
parser.parseLParen();
const n = parser.parseInteger();
numWarps.push(n);
parser.parseRParen();
const region = result.addRegion();
parser.parseRegion(region);
partitionIndex++;
}
result.addAttribute('numWarps', numWarps);
return true;
}
return super.parseOperation(parser, result);
}
};
_.MichelsonDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'michelson');
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
const typesWithParams = ['pair', 'list', 'option', 'or', 'map', 'big_map', 'set', 'contract', 'lambda'];
if (typesWithParams.includes(mnemonic)) {
parser.parseLess();
const types = [];
do {
types.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseGreater();
return new _.Type(`!${dialect}.${mnemonic}<${types.join(', ')}>`);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.PlanDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'plan');
this.registerCustomDirective('WithValuesTypes', this.parseWithValuesTypes.bind(this));
}
parseWithValuesTypes(parser, op, elementsRef, resultTypes, elementsTypes) {
const resultType = parser.parseType();
resultTypes.push(resultType);
if (op.types.length === 0) {
op.addTypes([resultType]);
} else {
op.types[0] = resultType;
}
// Derive element types from tensor's element type (matching C++ getElementTypeOrSelf)
const elementType = resultType.elementType || resultType;
// Populate element types for each element operand
if (elementsRef && elementsRef.length > 0) {
for (let i = 0; i < elementsRef.length; i++) {
elementsTypes.push(elementType);
}
}
}
};
_.KernelDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'kernel');
this.registerCustomDirective('KernelFunctionalType', this.parseKernelFunctionalType.bind(this));
}
parseKernelFunctionalType(parser, op /*, args */) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
parser.parseType();
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseArrow();
const resultTypes = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
resultTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
} else {
resultTypes.push(parser.parseType());
}
if (resultTypes.length > 0) {
op.addTypes(resultTypes);
}
}
inferResultTypes(op, vars) {
// kernel.combiner: result types are the first half of the input types
if (op.name.getStringRef() === 'kernel.combiner') {
const inputsEntry = vars.get('inputs');
if (inputsEntry?.types?.length > 0) {
const numPairs = Math.floor(inputsEntry.types.length / 2);
const resultTypes = inputsEntry.types.slice(0, numPairs);
op.addTypes(resultTypes);
return;
}
}
super.inferResultTypes(op, vars);
}
};
_.tensorrt = {};
_.tensorrt.TensorRTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tensorrt');
this.registerCustomAttribute('TensorRT_TopKOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ScatterModeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ResizeSelectorAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ResizeRoundModeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ResizeModeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ResizeCoordinateTransformationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ReduceOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_PaddingModeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_MatrixOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_LoopOutputAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_GatherModeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_FillOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ElementWiseOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_ActivationTypeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_UnaryOperationAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_TripLimitAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomAttribute('TensorRT_PoolingTypeAttr', this.parseEnumAttrBracket.bind(this));
this.registerCustomDirective('StaticIndexI64Array', this.parseStaticIndexI64Array.bind(this));
this.registerCustomDirective('StaticIndexI32Array', this.parseStaticIndexI32Array.bind(this));
}
parseEnumAttrBracket(parser) {
if (parser.parseOptionalLess()) {
const value = parser.parseKeyword();
parser.parseGreater();
return { value };
}
return null;
}
parseStaticIndexI64Array(parser, op, attrName = 'broadcast_dims') {
const values = [];
do {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
values.push(value);
} while (parser.parseOptionalComma());
op.addAttribute(attrName, values);
}
parseStaticIndexI32Array(parser, op, attrName = 'static_values') {
const values = [];
do {
const value = parser.parseOptionalInteger();
if (value === null) {
break;
}
values.push(value);
} while (parser.parseOptionalComma());
op.addAttribute(attrName, values);
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'tensorrt.for') {
return this.parseForOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseForOp(parser, result) {
const inductionVar = parser.parseOperand();
parser.parseEqual();
const unresolvedLb = parser.parseOperand();
parser.resolveOperand(unresolvedLb, null, result.operands);
parser.parseKeyword('to');
const unresolvedUb = parser.parseOperand();
parser.resolveOperand(unresolvedUb, null, result.operands);
parser.parseKeyword('step');
const unresolvedStep = parser.parseOperand();
parser.resolveOperand(unresolvedStep, null, result.operands);
parser.parseKeyword('init');
const regionArgs = [{ name: inductionVar.name, type: null }];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
let iterArgName = null;
const __iterArg = parser.parseOptionalOperand();
if (__iterArg) {
iterArgName = __iterArg.name;
}
if (parser.parseOptionalEqual()) {
const __initOp = parser.parseOptionalOperand();
if (__initOp) {
parser.resolveOperand(__initOp, null, result.operands);
if (iterArgName) {
regionArgs.push({ name: iterArgName, type: null });
}
}
}
parser.parseOptionalComma();
}
}
result.addTypes(parser.parseOptionalArrowTypeList());
parser.parseOptionalRegion(result.addRegion(), regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
};
_.executor = {};
_.executor.TableType = class extends _.Type {
constructor(body) {
super(null);
this.body = body;
}
toString() {
const bodyStr = this.body.map((t) => t.toString()).join(', ');
return `!executor.table<${bodyStr}>`;
}
};
_.executor.PointerType = class extends _.Type {
constructor(addressSpace) {
super(null);
this.addressSpace = addressSpace;
}
static parse(parser) {
parser.parseLess();
const addressSpace = parser.parseKeyword();
parser.parseGreater();
return new _.executor.PointerType(addressSpace);
}
toString() {
return `!executor.ptr<${this.addressSpace}>`;
}
};
_.executor.ExecutorFunctionType = class extends _.Type {
constructor(inputs, results, trailingVariadic) {
super(null);
this.inputs = inputs;
this.results = results;
this.trailingVariadic = trailingVariadic;
}
static parse(parser) {
parser.parseLess();
parser.parseLParen();
const inputs = [];
let trailingVariadic = false;
if (!parser.parseOptionalRParen()) {
if (parser.parseOptionalEllipsis()) {
trailingVariadic = true;
} else {
inputs.push(parser.parseType());
while (parser.parseOptionalComma()) {
if (parser.parseOptionalEllipsis()) {
trailingVariadic = true;
break;
}
inputs.push(parser.parseType());
}
}
parser.parseRParen();
}
parser.parseArrow();
parser.parseLParen();
const results = [];
if (!parser.parseOptionalRParen()) {
do {
results.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseGreater();
return new _.executor.ExecutorFunctionType(inputs, results, trailingVariadic);
}
toString() {
const input = this.inputs.map((t) => t.toString()).join(', ');
const result = this.results.map((t) => t.toString()).join(', ');
let variadic = '';
if (this.trailingVariadic) {
variadic = input ? ', ...' : '...';
}
return `!executor.func<(${input}${variadic}) -> (${result})>`;
}
};
_.executor.ExecutorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'executor');
this.registerCustomType('Executor_Table', this.parseTable.bind(this));
this.registerCustomDirective('ExecutorMixedIndices', this.parseExecutorMixedIndices.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'executor.func') {
parser.parseFunctionOp(result, true);
return true;
}
return super.parseOperation(parser, result);
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'executor.table.get') {
// ExtractTableValueOp::inferReturnTypes
const tableType = op.operands[0]?.type;
const index = op.attributes.get('index');
if (tableType instanceof _.executor.TableType && index !== undefined) {
const idx = typeof index === 'number' ? index : parseInt(index, 10);
if (idx >= 0 && idx < tableType.body.length) {
op.addTypes([tableType.body[idx]]);
return;
}
}
}
super.inferResultTypes(op, vars);
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
if (mnemonic === 'table') {
if (parser.parseOptionalLess()) {
const body = [];
if (!parser.parseOptionalGreater()) {
do {
const elementType = parser.parseType();
body.push(elementType);
} while (parser.parseOptionalComma());
parser.parseGreater();
}
return new _.executor.TableType(body);
}
return new _.executor.TableType([]);
}
if (mnemonic === 'ptr') {
return _.executor.PointerType.parse(parser);
}
if (mnemonic === 'func') {
return _.executor.ExecutorFunctionType.parse(parser);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parseTable(parser) {
if (parser.parseOptionalLess()) {
const body = [];
if (!parser.parseOptionalGreater()) {
do {
const elementType = parser.parseType();
body.push(elementType);
} while (parser.parseOptionalComma());
parser.parseGreater();
}
return new _.executor.TableType(body);
}
return null;
}
parseExecutorMixedIndices(parser, op /*, args */) {
const unresolvedOperands = [];
do {
const __mixedOp = parser.parseOptionalOperand();
if (__mixedOp) {
unresolvedOperands.push(__mixedOp);
} else {
parser.parseAttribute();
}
} while (parser.parseOptionalComma());
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, op.operands);
}
}
};
_.TFRTTestDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfrt_test');
this.registerCustomDirective('OptionalLoc', this.parseOptionalLoc.bind(this));
this.registerCustomDirective('DummyRegionRef', this.parseDummyRegionRef.bind(this));
this.registerCustomDirective('DummySuccessorRef', this.parseDummySuccessorRef.bind(this));
}
parseOptionalLoc(parser, op, attrName = 'loc') {
const loc = parser.parseOptionalLocationSpecifier();
if (loc) {
op.addAttribute(attrName, loc);
} else {
op.addAttribute(attrName, parser.getCurrentLocation());
}
}
parseDummyRegionRef() {
}
parseDummySuccessorRef() {
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
const opInfo = result.name.getRegisteredInfo();
if (!opInfo) {
return false;
}
if (opInfo.metadata.assemblyFormat === 'operands attr-dict') {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
for (const unresolved of unresolvedOperands) {
parser.resolveOperand(unresolved, null, result.operands);
}
this.inferResultTypes(result, new Map());
return true;
}
if (op === 'tfrt_test.do.async') {
const unresolvedOperands = parser.parseOperandList('none');
if (unresolvedOperands.length > 0) {
parser.resolveOperands(unresolvedOperands, unresolvedOperands.map(() => null), result.operands);
}
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type && type.results) {
type.results.forEach((resultType) => {
result.addTypes([resultType]);
});
}
}
parser.parseOptionalRegion(result.addRegion());
return true;
}
if (op === 'tfrt_test.benchmark') {
const __benchName = parser.parseOptionalString();
if (__benchName !== null) {
result.addAttribute('name', __benchName);
}
parser.parseLParen();
let __benchOp = parser.parseOptionalOperand();
while (__benchOp) {
let type = null;
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.resolveOperand(__benchOp, type, result.operands);
if (!parser.parseOptionalComma()) {
break;
}
__benchOp = parser.parseOptionalOperand();
}
parser.parseRParen();
for (let name = parser.parseOptionalKeyword(); name; name = parser.parseOptionalKeyword()) {
parser.parseEqual();
let value = null;
const __intVal = parser.parseOptionalInteger();
if (__intVal === null) {
const __strVal = parser.parseOptionalString();
if (__strVal === null) {
value = parser.parseKeyword();
} else {
value = __strVal;
}
} else {
value = __intVal;
}
result.addAttribute(name, value);
parser.parseOptionalComma();
}
parser.parseOptionalRegion(result.addRegion());
result.addTypes([new _.Type('!tfrt.chain')]);
return true;
}
return super.parseOperation(parser, result);
}
};
_.XeVMDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xevm');
this.registerCustomAttribute('XeVM_TruncfSrcElemTypeAttr', this.parseSrcElemTypeAttr.bind(this));
this.registerCustomAttribute('XeVM_TruncfDstElemTypeAttr', this.parseDstElemTypeAttr.bind(this));
this.registerCustomAttribute('XeVM_ExtfSrcElemTypeAttr', this.parseSrcElemTypeAttr.bind(this));
this.registerCustomAttribute('XeVM_ExtfDstElemTypeAttr', this.parseDstElemTypeAttr.bind(this));
}
parseSrcElemTypeAttr(parser) {
parser.parseKeyword('src_etype');
parser.parseEqual();
const value = parser.parseKeyword();
return new _.TypedAttr(value, null);
}
parseDstElemTypeAttr(parser) {
parser.parseKeyword('dst_etype');
parser.parseEqual();
const value = parser.parseKeyword();
return new _.TypedAttr(value, null);
}
};
_.VMVXDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'vmvx');
}
};
_.MLRTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'mlrt');
}
};
_.TFRTTensorDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfrt_tensor');
}
};
_.TFRTDHTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfrt_dht');
}
};
_.TFDDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'tfd');
}
};
_.ACCDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'acc');
this.registerCustomDirective('Var', this.parseVar.bind(this));
this.registerCustomDirective('AccVar', this.parseAccVar.bind(this));
this.registerCustomDirective('VarPtrType', this.parseVarPtrType.bind(this));
this.registerCustomDirective('DeviceTypeOperandsWithKeywordOnly', this.parseDeviceTypeOperandsWithKeywordOnly.bind(this));
this.registerCustomDirective('LoopControl', this.parseLoopControl.bind(this));
this.registerCustomDirective('WaitClause', this.parseWaitClause.bind(this));
this.registerCustomDirective('NumGangs', this.parseNumGangs.bind(this));
this.registerCustomDirective('DeviceTypeOperands', this.parseDeviceTypeOperands.bind(this));
this.registerCustomDirective('GangClause', this.parseGangClause.bind(this));
this.registerCustomDirective('CombinedConstructsLoop', this.parseCombinedConstructsLoop.bind(this));
this.registerCustomDirective('RecipeSym', this.parseRecipeSym.bind(this));
this.registerCustomDirective('OperandWithKeywordOnly', this.parseOperandWithKeywordOnly.bind(this));
this.registerCustomDirective('OperandsWithKeywordOnly', this.parseOperandsWithKeywordOnly.bind(this));
this.registerCustomDirective('DeviceTypeOperandsWithSegment', this.parseDeviceTypeOperandsWithSegment.bind(this));
this.registerCustomDirective('BindName', this.parseBindName.bind(this));
this.registerCustomDirective('RoutineGangClause', this.parseRoutineGangClause.bind(this));
this.registerCustomDirective('DeviceTypeArrayAttr', this.parseDeviceTypeArrayAttr.bind(this));
}
// custom<Var>($var) - receives ctx.get('var').operands
parseVar(parser, op, operands) {
if (!parser.parseOptionalKeyword('varPtr')) {
parser.parseKeyword('var');
}
parser.parseLParen();
operands.push(parser.parseOperand());
}
// custom<AccVar>($accVar, type($accVar)) - self-contained with type inline
parseAccVar(parser, op, operands, types) {
if (!parser.parseOptionalKeyword('accPtr')) {
parser.parseKeyword('accVar');
}
parser.parseLParen();
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
parser.parseRParen();
}
// custom<VarPtrType>(type($var), $varType) - receives ctx.get('var').types
parseVarPtrType(parser, op, types, varTypeAttrName) {
const type = parser.parseType();
types.push(type);
parser.parseRParen();
if (parser.parseOptionalKeyword('varType')) {
parser.parseLParen();
const varType = parser.parseType();
op.addAttribute(varTypeAttrName || 'varType', varType);
parser.parseRParen();
}
}
parseDeviceTypeOperandsWithKeywordOnly(parser, op, operands, types, deviceTypesVar, keywordOnlyVar) {
if (!parser.parseOptionalLParen()) {
op.addAttribute(keywordOnlyVar, [{ value: 'none' }]);
return;
}
const keywordOnlyAttrs = [];
let needComma = false;
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
keywordOnlyAttrs.push(parser.parseAttribute());
parser.parseOptionalComma();
}
needComma = true;
}
if (keywordOnlyAttrs.length > 0) {
op.addAttribute(keywordOnlyVar, keywordOnlyAttrs);
}
if (needComma) {
parser.parseComma();
}
const deviceTypes = [];
if (!parser.parseOptionalRParen()) {
parser.parseCommaSeparatedList('none', () => {
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
if (parser.parseOptionalLSquare()) {
deviceTypes.push(parser.parseAttribute());
parser.parseRSquare();
} else {
deviceTypes.push({ value: 'none' });
}
});
parser.parseRParen();
}
if (deviceTypes.length > 0) {
op.addAttribute(deviceTypesVar, deviceTypes);
}
}
parseLoopControl(parser, op, regionName, lowerbound, lowerboundTypes, upperbound, upperboundTypes, step, stepTypes) {
const inductionVars = [];
if (parser.parseOptionalKeyword('control')) {
parser.parseLParen();
while (!parser.parseOptionalRParen()) {
const value = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
inductionVars.push({ value, type });
parser.parseOptionalComma();
}
parser.parseEqual();
parser.parseLParen();
for (const o of parser.parseOperandList()) {
lowerbound.push(o);
}
for (const t of parser.parseColonTypeList()) {
lowerboundTypes.push(t);
}
parser.parseRParen();
parser.parseKeyword('to');
parser.parseLParen();
for (const o of parser.parseOperandList()) {
upperbound.push(o);
}
for (const t of parser.parseColonTypeList()) {
upperboundTypes.push(t);
}
parser.parseRParen();
parser.parseKeyword('step');
parser.parseLParen();
for (const o of parser.parseOperandList()) {
step.push(o);
}
for (const t of parser.parseColonTypeList()) {
stepTypes.push(t);
}
parser.parseRParen();
}
const region = op.addRegion();
parser.parseRegion(region, inductionVars);
}
parseWaitClause(parser, op, operands, types, deviceTypesAttr, segmentsAttr, hasDevNumAttr, keywordOnlyAttr) {
const keywordAttrs = [];
const deviceTypeAttrs = [];
const segments = [];
const devnum = [];
if (!parser.parseOptionalLParen()) {
keywordAttrs.push({ value: 'none' });
op.addAttribute(keywordOnlyAttr, keywordAttrs);
return;
}
let needComma = false;
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
keywordAttrs.push(parser.parseAttribute());
parser.parseOptionalComma();
}
needComma = true;
}
if (needComma) {
parser.parseComma();
}
do {
parser.parseLBrace();
const crtSize = operands.length;
if (parser.parseOptionalKeyword('devnum')) {
parser.parseColon();
devnum.push(true);
} else {
devnum.push(false);
}
parser.parseCommaSeparatedList('none', () => {
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
});
segments.push(operands.length - crtSize);
parser.parseRBrace();
if (parser.parseOptionalLSquare()) {
deviceTypeAttrs.push(parser.parseAttribute());
parser.parseRSquare();
} else {
deviceTypeAttrs.push({ value: 'none' });
}
} while (parser.parseOptionalComma());
parser.parseRParen();
if (deviceTypeAttrs.length > 0) {
op.addAttribute(deviceTypesAttr, deviceTypeAttrs);
}
if (keywordAttrs.length > 0) {
op.addAttribute(keywordOnlyAttr, keywordAttrs);
}
if (segments.length > 0) {
op.addAttribute(segmentsAttr, segments);
}
if (devnum.length > 0) {
op.addAttribute(hasDevNumAttr, devnum);
}
}
parseNumGangs(parser, op, operands, types, deviceTypesAttr, segmentsAttr) {
const deviceTypeAttrs = [];
const segments = [];
do {
parser.parseLBrace();
const crtSize = operands.length;
parser.parseCommaSeparatedList('none', () => {
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
});
segments.push(operands.length - crtSize);
parser.parseRBrace();
if (parser.parseOptionalLSquare()) {
deviceTypeAttrs.push(parser.parseAttribute());
parser.parseRSquare();
} else {
deviceTypeAttrs.push({ value: 'none' });
}
} while (parser.parseOptionalComma());
if (deviceTypeAttrs.length > 0) {
op.addAttribute(deviceTypesAttr, deviceTypeAttrs);
}
if (segments.length > 0) {
op.addAttribute(segmentsAttr, segments);
}
}
parseDeviceTypeOperands(parser, op, operands, types, deviceTypesAttr) {
const deviceTypeAttrs = [];
parser.parseCommaSeparatedList('none', () => {
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
if (parser.parseOptionalLSquare()) {
deviceTypeAttrs.push(parser.parseAttribute());
parser.parseRSquare();
} else {
deviceTypeAttrs.push({ value: 'none' });
}
});
if (deviceTypeAttrs.length > 0) {
op.addAttribute(deviceTypesAttr, deviceTypeAttrs);
}
}
parseGangClause(parser, op, gangOperands, gangTypes, gangArgTypeVar, gangDeviceTypeVar, gangSegmentsVar, gangOnlyVar) {
if (!parser.parseOptionalLParen()) {
op.addAttribute(gangOnlyVar, [{ value: 'none' }]);
return;
}
const gangOnlyAttrs = [];
let needComma = false;
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
const attr = parser.parseAttribute();
gangOnlyAttrs.push(attr);
parser.parseOptionalComma();
}
needComma = true;
}
if (gangOnlyAttrs.length > 0) {
op.addAttribute(gangOnlyVar, gangOnlyAttrs);
}
if (needComma) {
parser.parseOptionalComma();
}
const gangArgTypes = [];
const deviceTypes = [];
const segments = [];
while (parser.parseOptionalLBrace()) {
let segmentCount = 0;
do {
let argType = 'Num';
if (parser.parseOptionalKeyword('num')) {
parser.parseEqual();
argType = 'Num';
} else if (parser.parseOptionalKeyword('dim')) {
parser.parseEqual();
argType = 'Dim';
} else if (parser.parseOptionalKeyword('static')) {
parser.parseEqual();
argType = 'Static';
}
gangArgTypes.push({ value: argType });
const operand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(operand, type, op.operands);
segmentCount++;
} while (parser.parseOptionalComma());
parser.parseRBrace();
segments.push(segmentCount);
if (parser.parseOptionalLSquare()) {
const deviceType = parser.parseAttribute();
deviceTypes.push(deviceType);
parser.parseRSquare();
} else {
deviceTypes.push({ value: 'none' });
}
if (!parser.parseOptionalComma()) {
break;
}
}
parser.parseRParen();
if (gangArgTypes.length > 0) {
op.addAttribute('gangOperandsArgType', gangArgTypes);
}
if (deviceTypes.length > 0) {
op.addAttribute('gangOperandsDeviceType', deviceTypes);
}
if (segments.length > 0) {
op.addAttribute('gangOperandsSegments', segments);
}
}
parseCombinedConstructsLoop(parser, op, combinedAttr) {
const attr = parser.parseAttribute();
op.addAttribute(combinedAttr, attr);
}
parseRecipeSym(parser, op, recipeAttr) {
const attr = parser.parseAttribute();
op.addAttribute(recipeAttr, attr);
}
parseOperandWithKeywordOnly(parser, op, operands, types, asyncAttr) {
if (!parser.parseOptionalLParen()) {
op.addAttribute(asyncAttr, true);
return;
}
operands.push(parser.parseOperand());
parser.parseColon();
types.push(parser.parseType());
parser.parseRParen();
}
parseOperandsWithKeywordOnly(parser, op, operands, types, waitAttr) {
if (!parser.parseOptionalLParen()) {
op.addAttribute(waitAttr, true);
return;
}
parser.parseCommaSeparatedList('none', () => {
operands.push(parser.parseOperand());
});
parser.parseColon();
parser.parseCommaSeparatedList('none', () => {
types.push(parser.parseType());
});
parser.parseRParen();
}
parseDeviceTypeOperandsWithSegment(parser, op, operands, types, deviceTypesAttr, segmentsAttr) {
this.parseNumGangs(parser, op, operands, types, deviceTypesAttr, segmentsAttr);
}
parseBindName(parser, op, bindIdNameAttr, bindStrNameAttr, deviceIdTypesAttr, deviceStrTypesAttr) {
const bindIdNameAttrs = [];
const bindStrNameAttrs = [];
const deviceIdTypeAttrs = [];
const deviceStrTypeAttrs = [];
parser.parseCommaSeparatedList('none', () => {
const newAttr = parser.parseAttribute();
const isSymbolRef = typeof newAttr === 'object' && newAttr !== null && newAttr.type === 'SymbolRefAttr';
if (isSymbolRef) {
bindIdNameAttrs.push(newAttr);
} else {
bindStrNameAttrs.push(newAttr);
}
if (parser.parseOptionalLSquare()) {
const deviceType = parser.parseAttribute();
if (isSymbolRef) {
deviceIdTypeAttrs.push(deviceType);
} else {
deviceStrTypeAttrs.push(deviceType);
}
parser.parseRSquare();
} else if (isSymbolRef) {
deviceIdTypeAttrs.push({ value: 'none' });
} else {
deviceStrTypeAttrs.push({ value: 'none' });
}
});
if (bindIdNameAttrs.length > 0) {
op.addAttribute(bindIdNameAttr, bindIdNameAttrs);
}
if (bindStrNameAttrs.length > 0) {
op.addAttribute(bindStrNameAttr, bindStrNameAttrs);
}
if (deviceIdTypeAttrs.length > 0) {
op.addAttribute(deviceIdTypesAttr, deviceIdTypeAttrs);
}
if (deviceStrTypeAttrs.length > 0) {
op.addAttribute(deviceStrTypesAttr, deviceStrTypeAttrs);
}
}
parseRoutineGangClause(parser, op, gangAttr, gangDimAttr, gangDimDeviceTypeAttr) {
const gangAttrs = [];
const gangDimAttrs = [];
const gangDimDeviceTypeAttrs = [];
if (!parser.parseOptionalLParen()) {
gangAttrs.push({ value: 'none' });
op.addAttribute(gangAttr, gangAttrs);
return;
}
let needComma = false;
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
gangAttrs.push(parser.parseAttribute());
parser.parseOptionalComma();
}
needComma = true;
}
if (needComma) {
parser.parseComma();
}
parser.parseCommaSeparatedList('none', () => {
parser.parseKeyword('dim');
parser.parseColon();
gangDimAttrs.push(parser.parseAttribute());
if (parser.parseOptionalLSquare()) {
gangDimDeviceTypeAttrs.push(parser.parseAttribute());
parser.parseRSquare();
} else {
gangDimDeviceTypeAttrs.push({ value: 'none' });
}
});
parser.parseRParen();
if (gangAttrs.length > 0) {
op.addAttribute(gangAttr, gangAttrs);
}
if (gangDimAttrs.length > 0) {
op.addAttribute(gangDimAttr, gangDimAttrs);
}
if (gangDimDeviceTypeAttrs.length > 0) {
op.addAttribute(gangDimDeviceTypeAttr, gangDimDeviceTypeAttrs);
}
}
parseDeviceTypeArrayAttr(parser) {
if (parser.parseOptionalLSquare()) {
while (!parser.parseOptionalRSquare()) {
parser.parseAttribute();
parser.parseOptionalComma();
}
}
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'acc.compute_region') {
const entryArguments = [];
const launchOperands = [];
const inputOperands = [];
const inputTypes = [];
let streamOperand = null;
if (parser.parseOptionalKeyword('stream')) {
parser.parseLParen();
streamOperand = parser.parseOperand();
parser.parseColon();
const type = parser.parseType();
parser.resolveOperand(streamOperand, type, result.operands);
parser.parseRParen();
}
if (parser.parseOptionalKeyword('launch')) {
parser.parseLParen();
do {
const argName = parser.parseOperand();
parser.parseEqual();
const val = parser.parseOperand();
launchOperands.push(val);
entryArguments.push({ name: argName.name, type: new _.Type('!acc.par_width') });
} while (parser.parseOptionalComma());
parser.parseRParen();
}
if (parser.parseOptionalKeyword('ins')) {
parser.parseLParen();
do {
const argName = parser.parseOperand();
parser.parseEqual();
const val = parser.parseOperand();
inputOperands.push(val);
entryArguments.push({ name: argName.name, type: null });
} while (parser.parseOptionalComma());
parser.parseRParen();
parser.parseColon();
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
inputTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
let typeIdx = 0;
for (const arg of entryArguments) {
if (!arg.type) {
arg.type = inputTypes[typeIdx++];
}
}
}
const resultTypes = parser.parseOptionalArrowTypeList();
result.addTypes(resultTypes);
for (const op of launchOperands) {
parser.resolveOperand(op, null, result.operands);
}
for (let i = 0; i < inputOperands.length; i++) {
parser.resolveOperand(inputOperands[i], inputTypes[i] || null, result.operands);
}
const region = result.addRegion();
parser.parseRegion(region, entryArguments, true);
parser.parseOptionalAttrDict(result.attributes);
const segmentSizes = [launchOperands.length, inputOperands.length, streamOperand ? 1 : 0];
result.addAttribute('operandSegmentSizes', new _.DenseI32ArrayAttr(segmentSizes));
this.inferResultTypes(result, new Map());
return true;
}
return super.parseOperation(parser, result);
}
};
_.smt = {};
_.smt.BitVectorType = class extends _.Type {
constructor(width) {
super(null);
this.width = width;
}
toString() {
return `!smt.bv<${this.width}>`;
}
};
_.smt.SMTDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'smt');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'smt.eq' || op === 'smt.distinct') {
const unresolvedOperands = parser.parseOperandList();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const type = parser.parseType();
const types = unresolvedOperands.map(() => type);
parser.resolveOperands(unresolvedOperands, types, result.operands);
result.addTypes([new _.Type('!smt.bool')]);
return true;
}
if (op === 'smt.bv.repeat') {
const count = parser.parseInteger();
result.addAttribute('count', count);
parser.parseKeyword('times');
const unresolvedOperand = parser.parseOperand();
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const inputType = parser.parseType();
parser.resolveOperand(unresolvedOperand, inputType, result.operands);
// RepeatOp::parse - result width = input width * count
const inputWidth = inputType instanceof _.smt.BitVectorType ? inputType.width : 0;
const resultWidth = inputWidth * count;
result.addTypes([new _.smt.BitVectorType(resultWidth)]);
return true;
}
if (op === 'smt.int.constant') {
const value = parser.parseInteger();
result.addAttribute('value', value);
parser.parseOptionalAttrDict(result.attributes);
result.addTypes([new _.Type('!smt.int')]);
return true;
}
return super.parseOperation(parser, result);
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
switch (mnemonic) {
case 'bv': {
parser.parseLess();
const width = parser.parseInteger();
parser.parseGreater();
return new _.smt.BitVectorType(width);
}
case 'array': {
parser.parseLess();
parser.parseLSquare();
const domainType = parser.parseType();
parser.parseArrow();
const rangeType = parser.parseType();
parser.parseRSquare();
parser.parseGreater();
return new _.Type(`!${dialect}.array<[${domainType} -> ${rangeType}]>`);
}
case 'func': {
parser.parseLess();
parser.parseLParen();
const domainTypes = [];
if (!parser.parseOptionalRParen()) {
do {
domainTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRParen();
}
const rangeType = parser.parseType();
parser.parseGreater();
return new _.Type(`!${dialect}.func<(${domainTypes.join(', ')}) ${rangeType}>`);
}
case 'sort': {
parser.parseLess();
const identifier = parser.parseAttribute();
const params = [];
if (parser.parseOptionalLSquare()) {
do {
params.push(parser.parseType());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parser.parseGreater();
const paramStr = params.length > 0 ? `[${params.join(', ')}]` : '';
return new _.Type(`!${dialect}.sort<${identifier}${paramStr}>`);
}
default:
return new _.Type(`!${dialect}.${mnemonic}`);
}
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
inferResultTypes(op, vars) {
if (op.name.getStringRef() === 'smt.bv.concat') {
// ConcatOp::inferReturnTypes - result width = lhs width + rhs width
if (op.operands.length >= 2) {
const lhsType = op.operands[0].type;
const rhsType = op.operands[1].type;
const lhsWidth = lhsType instanceof _.smt.BitVectorType ? lhsType.width : 0;
const rhsWidth = rhsType instanceof _.smt.BitVectorType ? rhsType.width : 0;
const resultWidth = lhsWidth + rhsWidth;
op.addTypes([new _.smt.BitVectorType(resultWidth)]);
return;
}
}
super.inferResultTypes(op, vars);
}
};
_.MPMDDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'mpmd');
this.registerCustomType('mesh_tensor', this.parseMeshTensorType.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'mpmd.named_computation') {
return this.parseNamedComputationOp(parser, result);
}
if (op === 'mpmd.fragment') {
return this.parseFragmentOp(parser, result);
}
if (op === 'mpmd.fragment_call') {
return this.parseFragmentCallOp(parser, result);
}
if (op === 'mpmd.for') {
return this.parseForOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseNamedComputationOp(parser, result) {
// mpmd.named_computation<"name"(count)> (%inputs) (%block_args) { region } : (types) -> types
parser.parseLess();
// Parse single UserOriginAttr in short format: "name"(count)
const origin = this.parseUserOriginAttr(parser);
result.addAttribute('origin', origin);
parser.parseGreater();
const unresolvedInputs = parser.parseOperandList('paren');
const entryArguments = this.parseBlockArguments(parser);
parser.parseOptionalRegion(result.addRegion(), entryArguments);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedInputs, type.inputs, result.operands);
result.addTypes(type.results);
}
}
return true;
}
parseBlockArguments(parser) {
return parser.parseArgumentList('optionalParen', true);
}
parseFragmentOp(parser, result) {
// mpmd.fragment<mesh="m1", origin=["f1"], stage_id=N> (%inputs) {attrs} (%block_args) { region } : (types) -> type
parser.parseLess();
do {
const attrName = parser.parseKeyword();
parser.parseEqual();
// Use custom parser for origin attribute (array of UserOriginAttr)
const attrValue = attrName === 'origin' ? this.parseOriginArray(parser) : parser.parseAttribute();
result.addAttribute(attrName, attrValue);
} while (parser.parseOptionalComma());
parser.parseGreater();
const unresolvedInputs = parser.parseOperandList('paren');
parser.parseOptionalAttrDict(result.attributes);
const entryArguments = this.parseBlockArguments(parser);
parser.parseOptionalRegion(result.addRegion(), entryArguments);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedInputs, type.inputs, result.operands);
result.addTypes(type.results);
}
}
return true;
}
parseFragmentCallOp(parser, result) {
// mpmd.fragment_call<mesh="m1", origin=["f1"]> @callee(%args) {attrs} : (types) -> type
parser.parseLess();
do {
const attrName = parser.parseKeyword();
parser.parseEqual();
// Use custom parser for origin attribute (array of UserOriginAttr)
const attrValue = attrName === 'origin' ? this.parseOriginArray(parser) : parser.parseAttribute();
result.addAttribute(attrName, attrValue);
} while (parser.parseOptionalComma());
parser.parseGreater();
const callee = parser.parseOptionalSymbolName();
result.addAttribute('callee', callee);
const unresolvedArgs = parser.parseOperandList('paren');
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const type = parser.parseType();
if (type instanceof _.FunctionType) {
parser.resolveOperands(unresolvedArgs, type.inputs, result.operands);
result.addTypes(type.results);
}
}
return true;
}
parseForOp(parser, result) {
// Use parseOperandList for simple operand list (no types in syntax)
const inputs = parser.parseOperandList('paren');
const types = inputs.map(() => null);
parser.resolveOperands(inputs, types, result.operands);
parser.parseOptionalAttrDict(result.attributes);
const entryArguments = this.parseBlockArguments(parser);
parser.parseOptionalRegion(result.addRegion(), entryArguments);
if (parser.parseOptionalColon()) {
const resultTypes = [];
do {
resultTypes.push(parser.parseType());
} while (parser.parseOptionalComma());
result.addTypes(resultTypes);
}
return true;
}
parseMeshTensorType(parser, prefix) {
// Parse !mpmd.mesh_tensor<"mesh_name", tensor<shape>, sharding=<...>>
parser.parseLess();
const meshName = parser.parseString();
parser.parseComma();
const tensorType = parser.parseType();
const result = { name: prefix, meshName, tensorType };
if (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('sharding')) {
parser.parseEqual();
result.sharding = parser.parseAttribute();
}
}
parser.parseGreater();
return result;
}
// Parse UserOriginAttr in short format: "name"(count) where (count) is optional
parseUserOriginAttr(parser) {
const name = parser.parseString();
let transposeCount = 0;
if (parser.parseOptionalLParen()) {
const count = parser.parseInteger();
transposeCount = count;
parser.parseRParen();
}
return { name, transposeCount };
}
parseOriginArray(parser) {
const origins = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const origin = this.parseUserOriginAttr(parser);
origins.push(origin);
parser.parseOptionalComma();
}
return origins;
}
};
_.SdyDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'sdy');
this.registerCustomDirective('StrippedTensorShardingPerValueAttr', this.parseStrippedTensorShardingPerValueAttr.bind(this));
this.registerCustomDirective('SingleBlockRegionNoBlockId', this.parseSingleBlockRegionNoBlockId.bind(this));
this.registerCustomAttribute('Sdy_ListOfAxisRefLists', this.parseListOfAxisRefLists.bind(this));
this.registerCustomAttribute('Sdy_ManualAxes', this.parseManualAxes.bind(this));
this.registerCustomAttribute('Sdy_AllToAllParamList', this.parseAllToAllParamList.bind(this));
this.registerCustomAttribute('Sdy_TensorSharding', this.parseTensorShardingAttrWrap.bind(this));
this.registerCustomAttribute('Sdy_AxisRefList', this.parseAxisRefListWrap.bind(this));
this.registerCustomAttribute('Sdy_ReductionOp', this.parseEnumAttr.bind(this));
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'sdy.constant') {
return this.parseConstantOp(parser, result);
}
return super.parseOperation(parser, result);
}
parseStrippedTensorShardingPerValueAttr(parser, op, attrName) {
const shardings = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const sharding = this.parseTensorShardingAttr(parser);
shardings.push(sharding);
parser.parseOptionalComma();
}
if (attrName) {
op.addAttribute(attrName, shardings);
}
return shardings;
}
parseSingleBlockRegionNoBlockId(parser, op /*, args */) {
const entryArguments = [];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
const value = parser.parseOperand();
let type = null;
const attrs = [];
if (parser.parseOptionalColon()) {
type = parser.parseType();
}
parser.parseOptionalAttrDict(attrs);
entryArguments.push({ value, type, attributes: attrs.length > 0 ? attrs : undefined });
parser.parseOptionalComma();
}
}
const region = op.addRegion();
parser.parseRegion(region, entryArguments);
return region;
}
parseTensorShardingAttr(parser) {
parser.parseLess();
let meshOrRef = null;
meshOrRef = parser.parseOptionalSymbolName();
if (meshOrRef !== null) {
// meshOrRef is set
} else if (parser.parseOptionalKeyword('mesh')) {
meshOrRef = this.parseMeshAttr(parser);
} else {
parser.emitError("Expected '@' or 'mesh'");
}
parser.parseComma();
const dimShardings = this.parseDimensionShardings(parser);
let replicatedAxes = null;
let unreducedAxes = null;
while (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('replicated')) {
parser.parseEqual();
replicatedAxes = this.parseAxisRefList(parser);
} else if (parser.parseOptionalKeyword('unreduced')) {
parser.parseEqual();
unreducedAxes = this.parseAxisRefList(parser);
}
}
parser.parseGreater();
return { meshOrRef, dimShardings, replicatedAxes, unreducedAxes };
}
parseMeshAttr(parser) {
parser.parseLess();
const axes = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const name = parser.parseString();
parser.parseEqual();
const size = parser.parseInteger();
axes.push({ name, size });
parser.parseOptionalComma();
}
let deviceIds = null;
if (parser.parseOptionalComma()) {
if (parser.parseOptionalKeyword('device_ids')) {
parser.parseEqual();
deviceIds = parser.parseAttribute().value;
}
}
parser.parseGreater();
return { axes, deviceIds };
}
parseDimensionShardings(parser) {
const dimShardings = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const dimSharding = this.parseDimensionShardingAttr(parser);
dimShardings.push(dimSharding);
parser.parseOptionalComma();
}
return dimShardings;
}
parseDimensionShardingAttr(parser) {
const axes = [];
parser.parseLBrace();
let isClosed = true;
while (!parser.parseOptionalRBrace()) {
if (parser.parseOptionalQuestion()) {
isClosed = false;
parser.parseRBrace();
break;
}
const axis = this.parseAxisRefAttr(parser);
axes.push(axis);
if (parser.parseOptionalComma()) {
if (parser.parseOptionalQuestion()) {
isClosed = false;
parser.parseRBrace();
break;
}
}
}
let priority = null;
const priorityStr = parser.parseOptionalKeyword();
if (priorityStr && /^p\d+$/.test(priorityStr)) {
priority = parseInt(priorityStr.substring(1), 10);
}
return { axes, isClosed, priority };
}
parseAxisRefAttr(parser) {
const name = parser.parseString();
let subAxisInfo = null;
if (parser.parseOptionalColon()) {
subAxisInfo = this.parseSubAxisInfo(parser);
}
return { name, subAxisInfo };
}
parseSubAxisInfo(parser) {
parser.parseLParen();
const preSize = parser.parseInteger();
parser.parseRParen();
const size = parser.parseInteger();
return { preSize, size };
}
parseAxisRefList(parser) {
const axes = [];
parser.parseLBrace();
while (!parser.parseOptionalRBrace()) {
const axis = this.parseAxisRefAttr(parser);
axes.push(axis);
parser.parseOptionalComma();
}
return axes;
}
parseListOfAxisRefLists(parser) {
const lists = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const list = this.parseAxisRefList(parser);
lists.push(list);
parser.parseOptionalComma();
}
return { value: lists };
}
parseManualAxes(parser) {
const axes = [];
parser.parseLBrace();
while (!parser.parseOptionalRBrace()) {
const axis = parser.parseString();
axes.push(axis);
parser.parseOptionalComma();
}
return { value: axes };
}
parseAllToAllParamList(parser) {
const params = [];
parser.parseLSquare();
while (!parser.parseOptionalRSquare()) {
const param = this.parseAllToAllParam(parser);
params.push(param);
parser.parseOptionalComma();
}
return { value: params };
}
parseAllToAllParam(parser) {
const axes = this.parseAxisRefList(parser);
parser.parseColon();
const splitDim = parser.parseInteger();
parser.parseArrow();
const concatDim = parser.parseInteger();
return { axes, splitDim, concatDim };
}
parseTensorShardingAttrWrap(parser) {
return { value: this.parseTensorShardingAttr(parser) };
}
parseAxisRefListWrap(parser) {
return { value: this.parseAxisRefList(parser) };
}
parseConstantOp(parser, result) {
parser.parseOptionalAttrDict(result.attributes);
const attr = parser.parseAttribute();
result.addAttribute('value', attr.value || attr);
if (attr.type) {
if (result.types.length === 0) {
result.addTypes([attr.type]);
} else {
result.types[0] = attr.type;
}
}
return true;
}
};
_.XlaDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xla');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'xla.apply_indexing') {
const indexingMapAttr = parser.parseAttribute();
result.addAttribute('indexing_map_attr', indexingMapAttr);
const operands = [];
if (parser.parseOptionalLParen()) {
parser.parseOperands(operands);
parser.parseRParen();
}
if (parser.parseOptionalLSquare()) {
parser.parseOperands(operands);
parser.parseRSquare();
}
const indexType = new _.IndexType();
for (const operand of operands) {
parser.resolveOperand(operand, indexType, result.operands);
}
parser.parseOptionalAttrDict(result.attributes);
const map = indexingMapAttr.getIndexingMap().GetAffineMap();
const numResults = map.getNumResults();
for (let i = 0; i < numResults; i++) {
result.addTypes([indexType]);
}
return true;
}
// xla.loop (%dims)[%ivs] -> (%map_results) in #map iter_args(%args = %inits) -> (types) { body }
if (op === 'xla.loop') {
const unresolvedDims = [];
const unresolvedInits = [];
const regionArgs = [];
const indexType = new _.IndexType();
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const __dimOp = parser.parseOptionalOperand();
if (__dimOp) {
unresolvedDims.push(__dimOp);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
// Parse [%ivs] -> (%map_results) - these are block arguments
if (parser.parseOptionalLSquare()) {
if (!parser.parseOptionalRSquare()) {
do {
const __iv = parser.parseOptionalOperand();
if (__iv) {
regionArgs.push({ name: __iv.name, type: indexType });
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
}
if (parser.parseOptionalArrow()) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const __mapResult = parser.parseOptionalOperand();
if (__mapResult) {
regionArgs.push({ name: __mapResult.name, type: indexType });
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseKeyword('in');
const map = parser.parseAttribute();
result.addAttribute('indexing_map_attr', map);
const iterArgNames = [];
if (parser.parseOptionalKeyword('iter_args')) {
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const __iterArg = parser.parseOptionalOperand();
if (__iterArg) {
iterArgNames.push(__iterArg.name);
}
if (parser.parseOptionalEqual()) {
unresolvedInits.push(parser.parseOperand());
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
const resultTypes = [];
if (parser.parseOptionalArrow()) {
if (parser.parseOptionalLParen()) {
const types = parser.parseTypeList();
parser.parseRParen();
for (const t of types) {
result.addTypes([t]);
resultTypes.push(t);
}
} else {
const type = parser.parseType();
result.addTypes([type]);
resultTypes.push(type);
}
}
for (let i = 0; i < iterArgNames.length; i++) {
const argType = resultTypes[i] || null;
regionArgs.push({ name: iterArgNames[i], type: argType });
}
for (const dim of unresolvedDims) {
parser.resolveOperand(dim, indexType, result.operands);
}
for (let i = 0; i < unresolvedInits.length; i++) {
const initType = resultTypes[i] || null;
parser.resolveOperand(unresolvedInits[i], initType, result.operands);
}
parser.parseOptionalRegion(result.addRegion(), regionArgs);
parser.parseOptionalAttrDict(result.attributes);
return true;
}
return super.parseOperation(parser, result);
}
parseAttribute(parser, type) {
const keyword = parser.parseOptionalKeyword();
if (keyword === 'indexing_map') {
return _.IndexingMapAttr.parse(parser, type);
}
return null;
}
};
_.XlaGpuDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xla_gpu');
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// xla_gpu.shuffle_reduce(%ops) to N combiner=@func {attrs} : types
if (op === 'xla_gpu.shuffle_reduce') {
const unresolvedOperands = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const __srOp = parser.parseOptionalOperand();
if (__srOp) {
unresolvedOperands.push(__srOp);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseKeyword('to');
const maxDistance = parser.parseInteger();
result.addAttribute('max_distance', parseInt(maxDistance, 10));
parser.parseKeyword('combiner');
parser.parseEqual();
const combiner = parser.parseOptionalSymbolName();
result.addAttribute('combiner', new _.SymbolRefAttr(`@${combiner}`));
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const types = parser.parseTypeList();
parser.resolveOperands(unresolvedOperands, types, result.operands);
result.addTypes(types);
} else {
for (const operand of unresolvedOperands) {
parser.resolveOperand(operand, null, result.operands);
}
}
return true;
}
// xla_gpu.reduce (%inputs) inits(%inits) dimensions=[...] combiner=@func {attrs} : in_types to out_types
if (op === 'xla_gpu.reduce') {
const unresolvedInputs = [];
const unresolvedInits = [];
if (parser.parseOptionalLParen()) {
if (!parser.parseOptionalRParen()) {
do {
const __inp = parser.parseOptionalOperand();
if (__inp) {
unresolvedInputs.push(__inp);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
}
parser.parseKeyword('inits');
parser.parseLParen();
if (!parser.parseOptionalRParen()) {
do {
const __initOp = parser.parseOptionalOperand();
if (__initOp) {
unresolvedInits.push(__initOp);
}
} while (parser.parseOptionalComma());
parser.parseRParen();
}
parser.parseKeyword('dimensions');
parser.parseEqual();
const dimensions = parser.parseAttribute();
result.addAttribute('dimensions', dimensions);
parser.parseKeyword('combiner');
parser.parseEqual();
const combiner = parser.parseOptionalSymbolName();
result.addAttribute('combiner', new _.SymbolRefAttr(`@${combiner}`));
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalColon()) {
const inputTypes = parser.parseTypeList();
parser.resolveOperands(unresolvedInputs, inputTypes, result.operands);
parser.parseKeyword('to');
const outputTypes = parser.parseTypeList();
result.addTypes(outputTypes);
parser.resolveOperands(unresolvedInits, outputTypes, result.operands);
} else {
for (const input of unresolvedInputs) {
parser.resolveOperand(input, null, result.operands);
}
for (const init of unresolvedInits) {
parser.resolveOperand(init, null, result.operands);
}
}
return true;
}
return super.parseOperation(parser, result);
}
};
_.XTileLayoutAttr = class extends _.MemRefLayoutAttr {
constructor(minorToMajor) {
super();
this.minorToMajor = minorToMajor;
}
toString() {
return `#xtile.layout<[${this.minorToMajor.join(', ')}]>`;
}
};
_.XTileDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xtile');
}
parseAttribute(parser) {
const keyword = parser.parseOptionalKeyword();
if (keyword === 'layout') {
parser.parseLess();
parser.parseLSquare();
const values = [];
if (!parser.parseOptionalRSquare()) {
do {
values.push(parser.parseInteger());
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
parser.parseGreater();
return new _.XTileLayoutAttr(values);
}
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// xtile.entry_func - function-like op with custom format
if (op === 'xtile.entry_func') {
parser.parseFunctionOp(result, false);
return true;
}
return super.parseOperation(parser, result);
}
};
_.triton.xla.TritonXlaDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'triton_xla');
// Mark operations with custom directives as having custom assembly format
this.registerOperandName('triton_xla.extract', { name: 'triton_xla.extract', hasCustomAssemblyFormat: true });
this.registerOperandName('triton_xla.insert', { name: 'triton_xla.insert', hasCustomAssemblyFormat: true });
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
// triton_xla.extract from $src as memref<shape, layout> [offsets] [sizes] [strides] : result_type
if (op === 'triton_xla.extract') {
parser.parseKeyword('from');
const unresolvedSrc = parser.parseOperand();
parser.parseKeyword('as');
// Parse AsMemRefType: memref<shape, layout> - the memref type describes the pointer's layout
const memrefType = parser.parseType();
// Extract shape and layout from the memref type
this._extractMemRefInfo(memrefType, result, 'src_shape', 'src_layout');
this.parseDynamicIndexList(parser, result, 'offsets', 'static_offsets');
this.parseDynamicIndexList(parser, result, 'sizes', 'static_sizes');
this.parseDynamicIndexList(parser, result, 'strides', 'static_strides');
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const resultType = parser.parseType();
parser.resolveOperand(unresolvedSrc, memrefType, result.operands);
result.addTypes([resultType]);
return true;
}
// triton_xla.insert $src into $dst as memref<shape, layout> [offsets] [sizes] [strides] : src_type
if (op === 'triton_xla.insert') {
const unresolvedSrc = parser.parseOperand();
parser.parseKeyword('into');
const unresolvedDst = parser.parseOperand();
parser.parseKeyword('as');
// Parse AsMemRefType: memref<shape, layout> - the memref type describes the pointer's layout
const memrefType = parser.parseType();
// Extract shape and layout from the memref type
this._extractMemRefInfo(memrefType, result, 'dst_shape', 'dst_layout');
this.parseDynamicIndexList(parser, result, 'offsets', 'static_offsets');
this.parseDynamicIndexList(parser, result, 'sizes', 'static_sizes');
this.parseDynamicIndexList(parser, result, 'strides', 'static_strides');
parser.parseOptionalAttrDict(result.attributes);
parser.parseColon();
const srcType = parser.parseType();
parser.resolveOperand(unresolvedSrc, srcType, result.operands);
parser.resolveOperand(unresolvedDst, memrefType, result.operands);
return true;
}
return super.parseOperation(parser, result);
}
_extractMemRefInfo(memrefType, op, shapeAttrName, layoutAttrName) {
op.addAttribute(shapeAttrName, memrefType.shape.map((d) => Number(d)));
if (!(memrefType.layout instanceof _.XTileLayoutAttr)) {
throw new mlir.Error('Expected layout attribute');
}
op.addAttribute(layoutAttrName, memrefType.layout.minorToMajor.map((d) => Number(d)));
}
parseDynamicIndexList(parser, op, dynamicName, staticName) {
// Parse [val1, val2, ...] where vals can be %ssa or integer constants
parser.parseLSquare();
const staticValues = [];
if (!parser.parseOptionalRSquare()) {
do {
const __dynUnresolved = parser.parseOptionalOperand();
if (__dynUnresolved) {
parser.resolveOperand(__dynUnresolved, null, op.operands);
staticValues.push(-9223372036854775808n); // ShapedType::kDynamic
} else {
const __dynIntVal = parser.parseOptionalInteger('int64');
if (__dynIntVal !== null) {
staticValues.push(__dynIntVal);
}
}
} while (parser.parseOptionalComma());
parser.parseRSquare();
}
op.addAttribute(staticName, staticValues);
}
};
_.EnsembleDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'ensemble');
}
parseType(parser, dialect) {
const mnemonic = parser.parseKeyword();
const simpleTypes = ['physical_qubit', 'virtual_qubit', 'cbit', 'gate', 'gate_distribution', 'connectivity_graph'];
if (simpleTypes.includes(mnemonic)) {
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.PolyDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'poly');
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
if (mnemonic === 'poly') {
parser.parseLess();
const degreeBound = parser.parseInteger();
parser.parseGreater();
return new _.Type(`!${dialect}.poly<${degreeBound}>`);
}
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.NoisyDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'noisy');
}
parseType(parser, dialect) {
const mnemonic = parser.parseOptionalKeyword();
if (mnemonic) {
return new _.Type(`!${dialect}.${mnemonic}`);
}
parser.emitError(parser.getNameLoc(), `Unknown '${dialect}' type '${mnemonic}'`);
return null;
}
};
_.XtenNNDialect = class extends _.Dialect {
constructor(operations) {
super(operations, 'xten_nn');
}
parseAttribute(parser) {
const keyword = parser.parseOptionalKeyword();
if (keyword === 'dict_loc') {
parser.parseLParen();
const attrs = parser.parseAttribute();
parser.parseRParen();
return new _.XtenNNDictLoc(attrs);
}
return null;
}
parseOperation(parser, result) {
const op = result.name.getStringRef();
if (op === 'xten_nn.subgraph') {
return this.parseSubgraphOp(parser, result);
}
if (op === 'xten_nn.quantize' || op === 'xten_nn.dequantize') {
return this.parseQuantizeOp(parser, result);
}
if (op === 'xten_nn.kernel') {
return this.parseKernelOp(parser, result);
}
return super.parseOperation(parser, result);
}
// Parse: (%name = %value : type, ...) { body } -> result_types
// or: (%arg : type, ...) -> result_types (no body)
parseSubgraphOp(parser, result) {
const entryArgs = [];
if (parser.parseOptionalLParen()) {
while (!parser.parseOptionalRParen()) {
// Parse either (%name = %value : type) or (%value : type)
let blockArgName = null;
let operandRef = null;
const __firstOp = parser.parseOptionalOperand();
if (__firstOp) {
if (parser.parseOptionalEqual()) {
// (%name = %value : type) form
blockArgName = __firstOp;
operandRef = parser.parseOptionalOperand();
} else {
// (%value : type) form - value is both operand and block arg name
operandRef = __firstOp;
blockArgName = __firstOp;
}
}
parser.parseColon();
const type = parser.parseType();
if (operandRef) {
parser.resolveOperands([operandRef], [type], result.operands);
}
entryArgs.push({ value: blockArgName, type });
if (!parser.parseOptionalComma()) {
parser.parseOptionalRParen();
break;
}
}
}
parser.parseOptionalAttrDictWithKeyword(result.attributes);
parser.parseOptionalRegion(result.addRegion(), entryArgs, /* enableNameShadowing */ true);
if (parser.parseOptionalArrow()) {
result.addTypes(parser.parseTypeList());
}
return true;
}
// Parse: (%input: type) {attrs} -> type
parseQuantizeOp(parser, result) {
if (!parser.parseOptionalLParen()) {
return false;
}
const unresolvedOperands = [];
while (!parser.parseOptionalRParen()) {
const __qOp = parser.parseOptionalOperand();
if (__qOp) {
unresolvedOperands.push(__qOp);
}
parser.parseColon();
const type = parser.parseType();
parser.resolveOperands(unresolvedOperands, [type], result.operands);
unresolvedOperands.length = 0;
if (!parser.parseOptionalComma()) {
parser.parseOptionalRParen();
break;
}
}
parser.parseOptionalAttrDict(result.attributes);
result.addTypes(parser.parseArrowTypeList());
return true;
}
// Parse: "name" (%arg : type, ...) instantiation_args [...] {attrs} -> result_types
parseKernelOp(parser, result) {
const nameAttr = parser.parseAttribute();
result.addAttribute('name', nameAttr);
if (parser.parseOptionalLParen()) {
const unresolvedOperands = [];
const types = [];
while (!parser.parseOptionalRParen()) {
const __kOp = parser.parseOptionalOperand();
if (__kOp) {
unresolvedOperands.push(__kOp);
}
parser.parseColon();
types.push(parser.parseType());
if (!parser.parseOptionalComma()) {
parser.parseOptionalRParen();
break;
}
}
parser.resolveOperands(unresolvedOperands, types, result.operands);
}
if (parser.parseOptionalKeyword('instantiation_args')) {
const values = [];
const names = [];
this.parseKernelInstantiationArgs(parser, values, names);
result.addAttribute('instantiation_args', new _.ArrayAttr(values));
if (names.length > 0) {
result.addAttribute('instantiation_arg_names', new _.ArrayAttr(names));
}
}
parser.parseOptionalAttrDict(result.attributes);
if (parser.parseOptionalArrow()) {
result.addTypes(parser.parseTypeList());
}
return true;
}
parseKernelInstantiationArgs(parser, values, names) {
if (!parser.parseOptionalLSquare()) {
return;
}
while (!parser.parseOptionalRSquare()) {
const attr = parser.parseAttribute();
if (parser.parseOptionalEqual()) {
names.push(attr);
values.push(parser.parseAttribute());
} else {
values.push(attr);
}
parser.parseOptionalComma();
}
}
};
mlir.Metadata = class {
static async open(context) {
if (!mlir.Metadata._metadata) {
const data = await context.asset('mlir-metadata.json');
mlir.Metadata._metadata = new mlir.Metadata(data);
}
return mlir.Metadata._metadata;
}
constructor(data) {
this.operations = new Map();
if (data) {
const operations = JSON.parse(data);
for (const op of operations) {
const [dialectName] = op.name.split('.');
if (!this.operations.has(dialectName)) {
this.operations.set(dialectName, []);
}
this.operations.get(dialectName).push(op);
}
}
}
type(name) {
const [dialectName] = name.split('.');
const operations = this.operations.get(dialectName);
if (operations) {
const op = operations.find((op) => op.name === name);
if (op) {
return op;
}
}
return { name };
}
};
mlir.Error = class extends Error {
constructor(message) {
super(message);
this.name = 'Error loading MLIR model.';
}
};
export const ModelFactory = mlir.ModelFactory;