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
28386 lines
1.1 MiB
Plaintext
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;
|