diff --git a/mlir/utils/tree-sitter-mlir/.gitignore b/mlir/utils/tree-sitter-mlir/.gitignore new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/.gitignore @@ -0,0 +1,7 @@ +node_modules/ +src/ +bindings/ +binding.gyp +Cargo.toml +Cargo.lock +target/ diff --git a/mlir/utils/tree-sitter-mlir/README.md b/mlir/utils/tree-sitter-mlir/README.md --- a/mlir/utils/tree-sitter-mlir/README.md +++ b/mlir/utils/tree-sitter-mlir/README.md @@ -1,8 +1,18 @@ -tree-sitter-mlir -================ +# tree-sitter-mlir -Basic [tree-sitter](https://github.com/tree-sitter/tree-sitter) grammar for -MLIR following the [lang-ref](https://mlir.llvm.org/docs/LangRef/). +[tree-sitter](https://github.com/tree-sitter/tree-sitter) grammar for MLIR +following the [lang-ref](https://mlir.llvm.org/docs/LangRef/). The parser is +incomplete, and the bench statistics on the test files in the MLIR tree are as +follows: -Note: the directory in [LLVM repo](https://github.com/llvm/llvm-project/) -merely contains the grammar file(s) and not the NPM/generated code. +``` +Math, 100% passed +Builtin, 100% passed +Func, 100% passed +ControlFlow, 100% passed +Tensor, 93.33% passed +Arith, 83.33% passed +SCF, 88% passed +Affine, 73.08% passed +Linalg, 51.11% passed +``` diff --git a/mlir/utils/tree-sitter-mlir/bench.mjs b/mlir/utils/tree-sitter-mlir/bench.mjs new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/bench.mjs @@ -0,0 +1,45 @@ +'use strict'; + +import * as cp from "node:child_process"; +import * as process from "node:process"; +import * as path from "node:path"; +import * as os from "node:os"; +import glob from "glob"; + +const dialects = { + "Builtin": 100, + "Func": 100, + "Arith": 83, + "Math": 100, + "ControlFlow": 100, + "SCF": 88, + "Tensor": 93, + "Affine": 73, + "Linalg": 51, +}; + +let mlir_testdir = path.join(process.cwd(), "..", "..", "test", "Dialect"); + +function bench_dialect(dialect, min_pct) { + let testdir = path.join(mlir_testdir, dialect); + let testfiles = glob.sync(`${testdir}/*.mlir`, { ignore: [`${testdir}/invalid.mlir`] }); + let child = cp.spawn("npx", ["tree-sitter", "parse", "-q", "-s", ...testfiles], + { cwd: process.cwd() }); + let output = ""; + child.stdout.setEncoding("utf8"); + child.stdout.on("data", (data) => output += data); + child.on("close", () => { + let match = output.match(/success percentage: (\d+\.\d+)%/i); + let pass_pct = parseFloat(match[1]); + if (pass_pct < min_pct) { + console.log("%s, %f%% passed; minimum required is %d%%", dialect, pass_pct, min_pct); + process.exit(1); + } + console.log("%s, %f%% passed", dialect, pass_pct); + }); + child.on("error", (err) => console.log(err)); +} + +for (const [k, v] of Object.entries(dialects)) { + bench_dialect(k, v); +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/affine.js b/mlir/utils/tree-sitter-mlir/dialect/affine.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/affine.js @@ -0,0 +1,118 @@ +'use strict'; + +module.exports = { + affine_dialect : $ => prec.right(choice( + seq('affine.apply', + field('operand', + seq($.attribute, $._dim_and_symbol_use_list)), + field('attributes', optional($.attribute))), + + // operation ::= `affine.delinearize_index` $linear_index + // `into` ` ` + // `(` $basis `)` attr-dict `:` + // type($multi_index) + seq('affine.delinearlize_index', + field('operand', $.value_use), 'into', + field('basis', $._value_use_list_parens), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `affine.dma_start` ssa-use `[` + // multi-dim-affine-map-of-ssa-ids `]`, + // `[` multi-dim-affine-map-of-ssa-ids `]`, + // `[` multi-dim-affine-map-of-ssa-ids `]`, + // ssa-use `:` memref-type + seq(choice('affine.dma_start', 'affine.dma_wait'), + field('operands', + seq($.value_use, $._multi_dim_affine_expr_sq, + repeat(seq(',', $.value_use, + $._multi_dim_affine_expr_sq)))), + ',', field('numElements', $._value_use_list), + field('return', $._type_annotation)), + + // operation ::= `affine.for` ssa-id `=` lower-bound `to` + // upper-bound + // (`step` integer-literal)? `{` op* `}` + seq('affine.for', field('iv', $.value_use), '=', + field('lowerBound', + seq(optional(token('max')), $._bound)), + token('to'), + field('upperBound', + seq(optional(token('min')), $._bound)), + field('step', + optional(seq(token('step'), $.integer_literal))), + field('iter_args', + optional(seq(token('iter_args'), + $._value_assignment_list))), + field('return', optional($._function_return)), + field('body', $.region), + field('attributes', optional($.attribute))), + + // operation ::= `affine.if` if-op-cond `{` op* `}` + // (`else` `{` op* `}`)? if-op-cond ::= integer-set-attr + // dim-and-symbol-use-list + seq('affine.if', + field('condition', + seq($.attribute, $._dim_and_symbol_use_list)), + field('return', optional($._function_return)), + field('trueblk', $.region), + field('falseblk', + optional(seq(token('else'), $.region))), + field('attributes', optional($.attribute))), + + // operation ::= `affine.load` ssa-use `[` + // multi-dim-affine-map-of-ssa-ids `]` + // `:` memref-type + seq(choice('affine.load', 'affine.vector_load'), + field('operand', $.value_use), + field('multiDimAffineMap', + $._multi_dim_affine_expr_sq), + field('return', $._type_annotation)), + + // operation ::= `affine.min` affine-map-attribute + // dim-and-symbol-use-list + seq(choice('affine.min', 'affine.max'), + field('operand', + seq($.attribute, $._dim_and_symbol_use_list))), + + seq('affine.parallel', + field('iv', $._value_use_list_parens), '=', + field('lowerBound', $._multi_dim_affine_expr_parens), + token('to'), + field('upperBound', $._multi_dim_affine_expr_parens), + field('step', + optional(seq(token('step'), + $._multi_dim_affine_expr_parens))), + field('reduce', + optional(seq(token('reduce'), '(', + $.string_literal, + repeat(seq(',', $.string_literal)), + ')'))), + field('return', optional($._function_return)), + field('body', $.region)), + + seq('affine.prefetch', field('source', $.value_use), + field('indices', + optional($._multi_dim_affine_expr_sq)), + ',', field('isWrite', $.isWrite_attr), ',', + field('localityHint', $.localityHint_attr), ',', + field('isDataCache', $.isDataCache_attr), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `affine.store` ssa-use, ssa-use `[` + // multi-dim-affine-map-of-ssa-ids `]` + // `:` memref-type + seq(choice('affine.store', 'affine.vector_store'), + field('source', $.value_use), ',', + field('destination', $.value_use), + field('multiDimAffineMap', + $._multi_dim_affine_expr_sq), + field('return', $._type_annotation)), + + // operation ::= `affine.yield` attr-dict ($operands^ `:` + // type($operands))? + seq('affine.yield', + field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/arith.js b/mlir/utils/tree-sitter-mlir/dialect/arith.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/arith.js @@ -0,0 +1,142 @@ +'use strict'; + +module.exports = { + arith_dialect : $ => choice( + // operation ::= `arith.constant` attr-dict $value + seq('arith.constant', + field('attributes', optional($.attribute)), + field('value', $._literal_and_type)), + + // operation ::= `arith.addi` $lhs `,` $rhs attr-dict `:` + // type($result) operation ::= `arith.subi` $lhs `,` $rhs + // attr-dict `:` type($result) operation ::= `arith.divsi` + // $lhs `,` $rhs attr-dict `:` type($result) operation ::= + // `arith.divui` $lhs `,` $rhs attr-dict `:` type($result) + // operation ::= `arith.ceildivsi` $lhs `,` $rhs attr-dict + // `:` type($result) operation ::= `arith.ceildivui` $lhs + // `,` $rhs attr-dict `:` type($result) operation ::= + // `arith.floordivsi` $lhs `,` $rhs attr-dict `:` + // type($result) operation ::= `arith.remsi` $lhs `,` $rhs + // attr-dict `:` type($result) operation ::= `arith.remui` + // $lhs `,` $rhs attr-dict `:` type($result) operation ::= + // `arith.muli` $lhs `,` $rhs attr-dict `:` type($result) + // operation ::= `arith.mulsi_extended` $lhs `,` $rhs + // attr-dict `:` type($lhs) operation ::= + // `arith.mului_extended` $lhs `,` $rhs attr-dict `:` + // type($lhs) operation ::= `arith.andi` $lhs `,` $rhs + // attr-dict `:` type($result) operation ::= `arith.ori` + // $lhs `,` $rhs attr-dict `:` type($result) operation ::= + // `arith.xori` $lhs `,` $rhs attr-dict `:` type($result) + // operation ::= `arith.maxsi` $lhs `,` $rhs attr-dict `:` + // type($result) operation ::= `arith.maxui` $lhs `,` $rhs + // attr-dict `:` type($result) operation ::= `arith.minsi` + // $lhs `,` $rhs attr-dict `:` type($result) operation ::= + // `arith.minui` $lhs `,` $rhs attr-dict `:` type($result) + // operation ::= `arith.shli` $lhs `,` $rhs attr-dict `:` + // type($result) operation ::= `arith.shrsi` $lhs `,` $rhs + // attr-dict `:` type($result) operation ::= `arith.shrui` + // $lhs `,` $rhs attr-dict `:` type($result) + seq(choice('arith.addi', 'arith.subi', 'arith.divsi', + 'arith.divui', 'arith.ceildivsi', + 'arith.ceildivui', 'arith.floordivsi', + 'arith.remsi', 'arith.remui', 'arith.muli', + 'arith.mulsi_extended', 'arith.mului_extended', + 'arith.andi', 'arith.ori', 'arith.xori', + 'arith.maxsi', 'arith.maxui', 'arith.minsi', + 'arith.minui', 'arith.shli', 'arith.shrsi', + 'arith.shrui'), + field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `arith.addui_extended` $lhs `,` $rhs + // attr-dict `:` type($sum) + // `,` type($overflow) + seq('arith.addui_extended', field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('attributes', optional($.attribute)), + field('return', seq(':', $.type, ',', $.type))), + + // operation ::= `arith.addf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.divf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.maxf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.minf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.mulf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.remf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `arith.subf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + seq(choice('arith.addf', 'arith.divf', 'arith.maxf', + 'arith.minf', 'arith.mulf', 'arith.remf', + 'arith.subf'), + field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('fastmath', optional($.fastmath_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `arith.negf` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + seq(choice('arith.negf'), field('operand', $.value_use), + field('fastmath', optional($.fastmath_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `arith.cmpi` $predicate `,` $lhs `,` $rhs + // attr-dict `:` type($lhs) operation ::= `arith.cmpf` + // $predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs) + seq(choice('arith.cmpi', 'arith.cmpf'), + field('predicate', + choice('eq', 'ne', 'oeq', 'olt', 'ole', 'ogt', + 'oge', 'slt', 'sle', 'sgt', 'sge', 'ult', + 'ule', 'ugt', 'uge', $.string_literal)), + ',', field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `arith.extf` $in attr-dict `:` type($in) + // `to` type($out) operation ::= `arith.extsi` $in attr-dict + // `:` type($in) `to` type($out) operation ::= `arith.extui` + // $in attr-dict `:` type($in) `to` type($out) operation ::= + // `arith.fptosi` $in attr-dict `:` type($in) `to` + // type($out) operation ::= `arith.fptoui` $in attr-dict `:` + // type($in) `to` type($out) operation ::= + // `arith.index_cast` $in attr-dict `:` type($in) `to` + // type($out) operation ::= `arith.index_castui` $in + // attr-dict `:` type($in) `to` type($out) operation ::= + // `arith.sitofp` $in attr-dict `:` type($in) `to` + // type($out) operation ::= `arith.uitofp` $in attr-dict `:` + // type($in) `to` type($out) operation ::= `arith.bitcast` + // $in attr-dict `:` type($in) `to` type($out) operation ::= + // `arith.truncf` $in attr-dict `:` type($in) `to` + // type($out) operation ::= `arith.trunci` $in attr-dict `:` + // type($in) `to` type($out) + seq(choice('arith.extf', 'arith.extsi', 'arith.extui', + 'arith.fptosi', 'arith.fptoui', + 'arith.index_cast', 'arith.index_castui', + 'arith.sitofp', 'arith.uitofp', 'arith.bitcast', + 'arith.truncf', 'arith.trunci'), + field('in', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + seq('arith.select', field('cond', $.value_use), ',', + field('trueblk', $.value_use), ',', + field('falseblk', $.value_use), + field('return', $._type_annotation))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/bufferization.js b/mlir/utils/tree-sitter-mlir/dialect/bufferization.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/bufferization.js @@ -0,0 +1,32 @@ +'use strict'; + +module.exports = { + bufferization_dialect : $ => choice( + seq('bufferization.alloc_tensor', + field('in', $._value_use_list_parens), + field('copy', optional(seq(token('copy'), '(', + $.value_use, ')'))), + field('size_hint', + optional(seq(token('size_hint'), '=', + $.value_use))), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `bufferization.to_memref` $tensor + // attr-dict `:` type($memref) + seq('bufferization.to_memref', + field('tensor', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `bufferization.to_tensor` $memref + // (`restrict` $restrict^)? + // (`writable` $writable^)? attr-dict + // `:` type($memref) + seq('bufferization.to_tensor', + field('memref', $.value_use), + field('restrict', optional($.restrict_attr)), + field('writable', optional($.writable_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/builtin.js b/mlir/utils/tree-sitter-mlir/dialect/builtin.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/builtin.js @@ -0,0 +1,20 @@ +'use strict'; + +module.exports = { + builtin_dialect : $ => prec.right(choice( + // operation ::= `builtin.module` ($sym_name^)? + // attr-dict-with-keyword $bodyRegion + seq(choice('builtin.module', 'module'), + field('name', optional($.bare_id)), + field('attributes', optional($.attribute)), + field('body', $.region)), + + // operation ::= `builtin.unrealized_conversion_cast` + // ($inputs^ `:` type($inputs))? + // `to` type($outputs) attr-dict + seq(choice('builtin.unrealized_conversion_cast', + 'unrealized_conversion_cast'), + field('inputs', optional($._value_use_type_list)), + token('to'), field('outputs', $._type_list_no_parens), + field('attributes', optional($.attribute))))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/cf.js b/mlir/utils/tree-sitter-mlir/dialect/cf.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/cf.js @@ -0,0 +1,42 @@ +'use strict'; + +module.exports = { + cf_dialect : $ => prec.right(choice( + // operation ::= `cf.assert` $arg `,` $msg attr-dict + seq('cf.assert', field('argument', $.value_use), ',', + field('message', $.string_literal), + field('attributes', optional($.attribute))), + + // operation ::= `cf.br` $dest (`(` $destOperands^ `:` + // type($destOperands) `)`)? attr-dict + seq('cf.br', field('successor', $.successor), + field('attributes', optional($.attribute))), + + // operation ::= `cf.cond_br` $condition `,` + // $trueDest(`(` $trueDestOperands ^ `:` + // type($trueDestOperands)`)`)? `,` + // $falseDest(`(` $falseDestOperands ^ `:` + // type($falseDestOperands)`)`)? attr-dict + seq('cf.cond_br', field('condition', $.value_use), ',', + field('trueblk', $.successor), ',', + field('falseblk', $.successor), + field('attributes', optional($.attribute))), + + // operation ::= `cf.switch` $flag `:` type($flag) `,` `[` `\n` + // custom(ref(type($flag)),$defaultDestination, + // $defaultOperands, + // type($defaultOperands), + // $case_values, + // $caseDestinations, + // $caseOperands, + // type($caseOperands)) + // `]` + // attr-dict + seq('cf.switch', field('flag', $._value_use_and_type), ',', + '[', $.cf_case_label, $.successor, + repeat(seq(',', $.cf_case_label, $.successor)), ']', + field('attributes', optional($.attribute))), + )), + + cf_case_label : $ => seq(choice($.integer_literal, token('default')), ':') +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/func.js b/mlir/utils/tree-sitter-mlir/dialect/func.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/func.js @@ -0,0 +1,30 @@ +'use strict'; + +module.exports = { + func_dialect : $ => prec.right(choice( + // operation ::= `func.call_indirect` $callee `(` + // $callee_operands `)` attr-dict + // `:` type($callee) + // operation ::= `func.call` $callee `(` $operands `)` + // attr-dict + // `:` functional-type($operands, results) + seq(choice('func.call', 'call', 'func.call_indirect', + 'call_indirect'), + field('callee', $.symbol_ref_id), + field('operands', $._value_use_list_parens), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `func.constant` attr-dict $value `:` + // type(results) + seq(choice('func.constant', 'constant'), + field('attributes', optional($.attribute)), + field('value', $.symbol_ref_id), + field('return', $._function_type_annotation)), + + seq('func.func', $._op_func), + + seq(choice('func.return', 'return'), + field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/linalg.js b/mlir/utils/tree-sitter-mlir/dialect/linalg.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/linalg.js @@ -0,0 +1,80 @@ +'use strict'; + +module.exports = { + linalg_dialect : $ => prec.right(choice( + seq(choice( + 'linalg.batch_matmul', + 'linalg.batch_matmul_transpose_b', + 'linalg.batch_matvec', + 'linalg.batch_reduce_matmul', 'linalg.broadcast', + 'linalg.conv_1d_ncw_fcw', 'linalg.conv_1d_nwc_wcf', + 'linalg.conv_1d', 'linalg.conv_2d_nchw_fchw', + 'linalg.conv_2d_ngchw_fgchw', + 'linalg.conv_2d_nhwc_fhwc', + 'linalg.conv_2d_nhwc_hwcf', + 'linalg.conv_2d_nhwc_hwcf_q', 'linalg.conv_2d', + 'linalg.conv_3d_ndhwc_dhwcf', + 'linalg.conv_3d_ndhwc_dhwcf_q', 'linalg.conv_3d', + 'linalg.copy', 'linalg.depthwise_conv_1d_nwc_wc', + 'linalg.depthwise_conv_1d_nwc_wcm', + 'linalg.depthwise_conv_2d_nchw_chw', + 'linalg.depthwise_conv_2d_nhwc_hwc', + 'linalg.depthwise_conv_2d_nhwc_hwc_q', + 'linalg.depthwise_conv_2d_nhwc_hwcm', + 'linalg.depthwise_conv_2d_nhwc_hwcm_q', + 'linalg.depthwise_conv_3d_ndhwc_dhwc', + 'linalg.depthwise_conv_3d_ndhwc_dhwcm', + 'linalg.dot', 'linalg.elemwise_binary', + 'linalg.elemwise_unary', 'linalg.fill', + 'linalg.fill_rng_2d', 'linalg.matmul', + 'linalg.matmul_transpose_b', + 'linalg.matmul_unsigned', 'linalg.matvec', + 'linalg.mmt4d', 'linalg.pooling_nchw_max', + 'linalg.pooling_nchw_sum', + 'linalg.pooling_ncw_max', 'linalg.pooling_ncw_sum', + 'linalg.pooling_ndhwc_max', + 'linalg.pooling_ndhwc_min', + 'linalg.pooling_ndhwc_sum', + 'linalg.pooling_nhwc_max', + 'linalg.pooling_nhwc_max_unsigned', + 'linalg.pooling_nhwc_min', + 'linalg.pooling_nhwc_min_unsigned', + 'linalg.pooling_nhwc_sum', + 'linalg.pooling_nwc_max', + 'linalg.pooling_nwc_max_unsigned', + 'linalg.pooling_nwc_min', + 'linalg.pooling_nwc_min_unsigned', + 'linalg.pooling_nwc_sum', + 'linalg.quantized_batch_matmul', + 'linalg.quantized_matmul', 'linalg.vecmat'), + repeat1($._ins_outs_attributes), + field('return', optional($._function_return))), + + seq('linalg.generic', repeat1($._ins_outs_attributes), + field('body', $.region), + field('return', optional($._function_return))), + + // operation ::= `linalg.index` $dim attr-dict `:` + // type($result) + seq('linalg.index', field('dimension', $.integer_literal), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + seq(choice('linalg.map', 'linalg.reduce'), + repeat1($._ins_outs_attributes), + field('arguments', $.block_arg_list), + field('body', $.region), + field('return', optional($._function_return))), + + seq('linalg.yield', + field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))))), + + _ins_outs_attributes : $ => choice($._ins, $._outs, $.attribute, + $._attribute_entry_list), + _ins : $ => seq(token('ins'), '(', $._value_use_type_list, ')'), + _outs : $ => seq(token('outs'), '(', $._value_use_type_list, ')'), + _attribute_entry_list : $ => seq($.bare_attribute_entry, + repeat(seq(',', $.bare_attribute_entry))), + bare_attribute_entry : $ => seq($.bare_id, '=', $.attribute_value) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/llvm.js b/mlir/utils/tree-sitter-mlir/dialect/llvm.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/llvm.js @@ -0,0 +1,10 @@ +'use strict'; + +module.exports = { + llvm_dialect : $ => prec.right(choice( + seq('llvm.func', $._op_func), + + seq('llvm.return', + field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/math.js b/mlir/utils/tree-sitter-mlir/dialect/math.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/math.js @@ -0,0 +1,129 @@ +'use strict'; + +module.exports = { + math_dialect : $ => choice( + // operation ::= `math.absf` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.atan` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.cbrt` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.ceil` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.cos` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.erf` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.exp` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.exp2` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.expm1` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.floor` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.log` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.log10` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.log1p` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.log2` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.round` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.roundeven` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.rsqrt` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.sin` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.sqrt` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.tan` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.tanh` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.trunc` $operand (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + seq(choice('math.absf', 'math.atan', 'math.cbrt', + 'math.ceil', 'math.cos', 'math.erf', 'math.exp', + 'math.exp2', 'math.expm1', 'math.floor', + 'math.log', 'math.log10', 'math.log1p', + 'math.log2', 'math.round', 'math.roundeven', + 'math.rsqrt', 'math.sin', 'math.sqrt', 'math.tan', + 'math.tanh', 'math.trunc'), + field('operand', $.value_use), + field('fastmath', optional($.fastmath_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `math.absi` $operand attr-dict `:` + // type($result) operation ::= `math.ctlz` $operand attr-dict + // `:` type($result) operation ::= `math.cttz` $operand + // attr-dict `:` type($result) operation ::= `math.ctpop` + // $operand attr-dict `:` type($result) + seq(choice('math.absi', 'math.ctlz', 'math.cttz', + 'math.ctpop'), + field('operand', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `math.atan2` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.copysign` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + // operation ::= `math.fpowi` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($lhs) `,` type($rhs) + // operation ::= `math.powf` $lhs `,` $rhs (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + seq(choice('math.atan2', 'math.copysign', 'math.fpowi', + 'math.powf'), + field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('fastmath', optional($.fastmath_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `math.ipowi` $lhs `,` $rhs attr-dict `:` + // type($result) + seq('math.ipowi', field('lhs', $.value_use), ',', + field('rhs', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `math.fma` $a `,` $b `,` $c (`fastmath` `` + // $fastmath^)? + // attr-dict `:` type($result) + seq('math.fma', field('a', $.value_use), ',', + field('b', $.value_use), ',', field('c', $.value_use), + field('fastmath', optional($.fastmath_attr)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/memref.js b/mlir/utils/tree-sitter-mlir/dialect/memref.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/memref.js @@ -0,0 +1,157 @@ +'use strict'; + +module.exports = { + memref_dialect : $ => choice( + // operation ::= `memref.assume_alignment` $memref `,` + // $alignment attr-dict `:` type($memref) + seq('memref.assume_alignment', + field('memref', $.value_use), ',', + field('alignment', $.integer_literal), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.alloc` `(`$dynamicSizes`)` (`` `[` + // $symbolOperands^ `]`)? attr-dict + // `:` type($memref) + // operation ::= `memref.alloca` `(`$dynamicSizes`)` (`` + // `[` $symbolOperands^ `]`)? attr-dict + // `:` type($memref) + seq(choice('memref.alloc', 'memref.alloca'), + field('dynamicSizes', $._value_use_list_parens), + field('symbolOperands', optional($._dense_idx_list)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.cast` $source attr-dict `:` + // type($source) `to` type($dest) + seq('memref.cast', field('in', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.copy` $source `,` $target + // attr-dict + // `:` type($source) `to` type($target) + seq('memref.copy', field('source', $.value_use), ',', + field('target', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.collapse_shape` $src + // $reassociation attr-dict + // `:` type($src) `into` type($result) + // operation ::= `memref.expand_shape` $src $reassociation + // attr-dict + // `:` type($src) `into` type($result) + seq(choice('memref.collapse_shape', 'memref.expand_shape'), + field('source', $.value_use), + field('reassociation', $.nested_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.dealloc` $memref attr-dict `:` + // type($memref) + seq('memref.dealloc', field('memref', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.dim` attr-dict $source `,` $index + // `:` type($source) + seq('memref.dim', + field('attributes', optional($.attribute)), + field('source', $.value_use), ',', + field('index', $.value_use), + field('return', $._type_annotation)), + + // operation ::= `memref.load` $memref `[` $indices `]` + // attr-dict `:` type($memref) + seq('memref.load', + field('memref', seq($.value_use, $._dense_idx_list)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + seq('memref.prefetch', field('source', $.value_use), + field('indices', optional($._dense_idx_list)), ',', + field('isWrite', $.isWrite_attr), ',', + field('localityHint', $.localityHint_attr), ',', + field('isDataCache', $.isDataCache_attr), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.rank` $memref attr-dict `:` + // type($memref) + seq('memref.rank', field('memref', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.realloc` $source (`(` + // $dynamicResultSize^ `)`)? attr-dict + // `:` type($source) `to` type(results) + seq('memref.realloc', field('source', $.value_use), + field('dynamicResultSize', + optional($._value_use_list_parens)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.reinterpret_cast` $source `to` + // `offset` `` `:` + // custom($offsets, + // $static_offsets) + // `` `,` `sizes` `` `:` + // custom($sizes, + // $static_sizes) + // `` `,` `strides` `` `:` + // custom($strides, + // $static_strides) attr-dict `:` + // type($source) `to` type($result) + seq('memref.reinterpret_cast', + field('source', $.value_use), token('to'), + field('offset', seq(token('offset'), ':', + $._dense_idx_list, ',')), + field('sizes', seq(token('sizes'), ':', + $._dense_idx_list, ',')), + field('strides', + seq(token('strides'), ':', $._dense_idx_list)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.reshape` $source `(` $shape `)` + // attr-dict + // `:` functional-type(operands, results) + seq('memref.reshape', field('source', $.value_use), + field('shape', seq('(', $.value_use, ')')), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `memref.store` $value `,` $memref `[` + // $indices `]` attr-dict + // `:` type($memref) + seq('memref.store', field('source', $.value_use), ',', + field('destination', $.value_use), + field('indices', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.subview` $source `` + // custom($offsets, + // $static_offsets) + // custom($sizes, + // $static_sizes) + // custom($strides, + // $static_strides) attr-dict `:` + // type($source) `to` type($result) + seq('memref.subview', field('source', $.value_use), + field('offsets', $._dense_idx_list), + field('sizes', $._dense_idx_list), + field('strides', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `memref.view` $source `[` $byte_shift `]` + // `` `[` $sizes `]` attr-dict + // `:` type($source) `to` type(results) + seq('memref.view', field('source', $.value_use), + field('byte_shift', $._dense_idx_list), + field('sizes', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/scf.js b/mlir/utils/tree-sitter-mlir/dialect/scf.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/scf.js @@ -0,0 +1,107 @@ +'use strict'; + +module.exports = { + scf_dialect : $ => prec.right(choice( + // operation ::= `scf.condition` `(` $condition `)` attr-dict + // ($args^ `:` type($args))? + seq('scf.condition', + field('condition', $._value_use_list_parens), + field('attributes', optional($.attribute)), + field('arguments', optional($._value_use_type_list))), + + seq('scf.execute_region', + field('return', optional($._function_return)), + field('body', $.region), + field('attributes', optional($.attribute))), + + seq('scf.if', field('condition', $.value_use), + field('return', optional($._function_return)), + field('trueblk', $.region), + field('falseblk', + optional(seq(token('else'), $.region)))), + + // operation ::= `scf.index_switch` $arg attr-dict (`->` + // type($results)^)? + // custom($cases, $caseRegions) + // `\n` + // `` `default` $defaultRegion + seq('scf.index_switch', field('flag', $._value_use_and_type), + field('attributes', optional($.attribute)), + field('result', optional($._function_return)), + $.scf_case_label, $.region, + repeat(seq($.scf_case_label, $.region))), + + // scf.for %iv = %lb to %ub step %step { + // ... // body + // } + seq('scf.for', field('iv', $.value_use), '=', + field('lb', $.value_use), token('to'), + field('ub', $.value_use), + field('step', seq(token('step'), $.value_use)), + field('iter_args', + optional(seq(token('iter_args'), + $._value_assignment_list))), + field('return', optional($._function_return)), + field('body', $.region), + field('attributes', optional($.attribute))), + + seq('scf.forall', field('iv', $._value_use_list_parens), + field('bounds', + seq(choice(seq('=', $._value_use_list_parens, + token('to')), + token('in')), + $._value_use_list_parens)), + field('step', optional(seq(token('step'), + $._value_use_list_parens))), + field('shared_outs', + optional(seq(token('shared_outs'), + $._value_assignment_list))), + field('return', optional($._function_return)), + field('body', $.region), + field('attributes', optional($.attribute))), + + seq('scf.forall.in_parallel', field('body', $.region), + field('attributes', optional($.attribute))), + + seq('scf.parallel', field('iv', $._value_use_list_parens), + '=', field('lb', $._value_use_list_parens), token('to'), + field('ub', $._value_use_list_parens), + field('step', + seq(token('step'), $._value_use_list_parens)), + field('init', optional(seq(token('init'), + $._value_use_list_parens))), + field('return', optional($._function_return)), + field('body', $.region), + field('attributes', optional($.attribute))), + + seq('scf.reduce', field('operand', $._value_use_list_parens), + field('return', $._type_annotation), + field('body', $.region)), + + // operation ::= `scf.reduce.return` $result attr-dict `:` + // type($result) + seq('scf.reduce.return', field('result', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // op ::= `scf.while` assignments `:` function-type region + // `do` region + // `attributes` attribute-dict + // initializer ::= /* empty */ | `(` assignment-list `)` + seq('scf.while', + field('assignments', optional($._value_assignment_list)), + field('return', $._function_type_annotation), + field('condblk', $.region), 'do', + field('doblk', $.region), + field('attributes', + optional(seq('attributes', $.attribute)))), + + // operation ::= `scf.yield` attr-dict ($results^ `:` + // type($results))? + seq('scf.yield', field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))), + )), + + scf_case_label : $ => choice(seq(token('case'), $.integer_literal), + token('default')) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/tensor.js b/mlir/utils/tree-sitter-mlir/dialect/tensor.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/tensor.js @@ -0,0 +1,201 @@ +'use strict'; + +module.exports = { + tensor_dialect : $ => choice( + // operation ::= `tensor.empty` `(`$dynamicSizes`)` + // attr-dict `:` type($result) + seq('tensor.empty', + field('dynamicSizes', $._value_use_list_parens), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.cast` $source attr-dict `:` + // type($source) `to` type($dest) + seq('tensor.cast', field('in', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.dim` attr-dict $source `,` $index + // `:` type($source) + seq('tensor.dim', + field('attributes', optional($.attribute)), + field('tensor', $.value_use), ',', + field('index', $.value_use), + field('return', $._type_annotation)), + + // operation ::= `tensor.collapse_shape` $src + // $reassociation attr-dict `:` type($src) + // `into` type($result) + seq(choice('tensor.collapse_shape', 'tensor.expand_shape'), + field('tensor', $.value_use), + field('reassociation', $.nested_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.extract` $tensor `[` $indices `]` + // attr-dict `:` type($tensor) + seq('tensor.extract', field('tensor', $.value_use), + field('indices', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.insert` $scalar `into` $dest `[` + // $indices `]` attr-dict + // `:` type($dest) + seq('tensor.insert', field('scalar', $.value_use), + token('into'), field('destination', $.value_use), + field('indices', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.extract_slice` $source `` + // custom($offsets, + // $static_offsets) + // custom($sizes, + // $static_sizes) + // custom($strides, + // $static_strides) attr-dict `:` + // type($source) `to` type($result) + seq('tensor.extract_slice', field('tensor', $.value_use), + field('offsets', $._dense_idx_list), + field('sizes', $._dense_idx_list), + field('strides', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.insert_slice` $source `into` $dest + // `` + // custom($offsets, + // $static_offsets) + // custom($sizes, + // $static_sizes) + // custom($strides, + // $static_strides) attr-dict `:` + // type($source) `into` type($dest) + // operation ::= `tensor.parallel_insert_slice` $source + // `into` $dest `` + // custom($offsets, + // $static_offsets) + // custom($sizes, + // $static_sizes) + // custom($strides, + // $static_strides) attr-dict `:` + // type($source) `into` type($dest) + seq(choice('tensor.insert_slice', + 'tensor.parallel_insert_slice'), + field('source', $.value_use), token('into'), + field('destination', $.value_use), + field('offsets', $._dense_idx_list), + field('sizes', $._dense_idx_list), + field('strides', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.from_elements` $elements attr-dict + // `:` type($result) + seq('tensor.from_elements', + field('elements', optional($._value_use_list)), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.gather` $source `[` $indices `]` + // `gather_dims` `(` $gather_dims `)` + // (`unique` $unique^)? + // attr-dict + // `:` functional-type(operands, results) + seq('tensor.gather', field('source', $.value_use), + field('indices', $._dense_idx_list), + field('gatherDims', $.gather_dims_attr), + field('unique', optional($.unique_attr)), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `tensor.scatter` $source `into` $dest `[` + // $indices `]` + // `scatter_dims` `(` $scatter_dims `)` + // (`unique` $unique^)? + // attr-dict + // `:` functional-type(operands, results) + seq('tensor.scatter', field('source', $.value_use), + token('into'), field('destination', $.value_use), + field('indices', $._dense_idx_list), + field('scatterDims', $.scatter_dims_attr), + field('unique', optional($.unique_attr)), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `tensor.pad` $source + // (`nofold` $nofold^)? + // `low` `` custom($low, + // $static_low) `high` `` + // custom($high, + // $static_high) $region attr-dict `:` + // type($source) `to` type($result) + seq('tensor.pad', field('source', $.value_use), + field('nofold', optional($.nofold_attr)), + field('low', seq(token('low'), $._dense_idx_list)), + field('high', seq(token('high'), $._dense_idx_list)), + field('body', $.region), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.reshape` $source `(` $shape `)` + // attr-dict + // `:` functional-type(operands, results) + seq('tensor.reshape', field('tensor', $.value_use), + field('shape', $._value_use_list_parens), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `tensor.splat` $input attr-dict `:` + // type($aggregate) + seq('tensor.splat', field('input', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.pack` $source + // (`padding_value` `(` $padding_value^ `:` + // type($padding_value) `)`)? + // (`outer_dims_perm` `=` $outer_dims_perm^)? + // `inner_dims_pos` `=` $inner_dims_pos + // `inner_tiles` `=` + // custom($inner_tiles, + // $static_inner_tiles) `into` $dest + // attr-dict `:` type($source) `->` + // type($dest) + // operation ::= `tensor.unpack` $source + // (`outer_dims_perm` `=` $outer_dims_perm^)? + // `inner_dims_pos` `=` $inner_dims_pos + // `inner_tiles` `=` + // custom($inner_tiles, + // $static_inner_tiles) `into` $dest + // attr-dict `:` type($source) `->` + // type($dest) + seq(choice('tensor.pack', 'tensor.unpack'), + field('source', $.value_use), + field('padding_value', + optional(seq(token('padding_value'), '(', + $._value_use_and_type, ')'))), + field('outer_dims_perm', + optional($.outer_dims_perm_attr)), + field('inner_dims_pos', $.inner_dims_pos_attr), + field('inner_tiles', $.inner_tiles_attr), + token('into'), field('destination', $.value_use), + field('return', $._function_type_annotation)), + + // operation ::= `tensor.generate` $dynamicExtents $body + // attr-dict `:` type($result) + seq('tensor.generate', + field('dynamicExtents', $._value_use_list), + field('body', $.region), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `tensor.rank` $tensor attr-dict `:` + // type($tensor) operation ::= `tensor.yield` $value + // attr-dict `:` type($value) + seq(choice('tensor.rank', 'tensor.yield'), + field('tensor', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation))) +} diff --git a/mlir/utils/tree-sitter-mlir/dialect/vector.js b/mlir/utils/tree-sitter-mlir/dialect/vector.js new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/dialect/vector.js @@ -0,0 +1,121 @@ +'use strict'; + +module.exports = { + vector_dialect : $ => prec.right(choice( + // operation ::= `vector.bitcast` $source attr-dict `:` + // type($source) `to` type($result) operation ::= + // `vector.broadcast` $source attr-dict `:` type($source) + // `to` type($vector) operation ::= `vector.shape_cast` + // $source attr-dict `:` type($source) `to` type($result) + // operation ::= `vector.type_cast` $memref attr-dict `:` + // type($memref) `to` type($result) + seq(choice('vector.bitcast', 'vector.broadcast', + 'vector.shape_cast', 'vector.type_cast'), + field('in', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.constant_mask` $mask_dim_sizes + // attr-dict `:` type(results) + seq('vector.constant_mask', + field('mask', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.create_mask` $operands attr-dict + // `:` type(results) + seq('vector.create_mask', + field('operands', $._value_use_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.extract` $vector `` $position + // attr-dict `:` type($vector) operation ::= `vector.load` + // $base `[` $indices `]` attr-dict + // `:` type($base) `,` type($nresult) + // operation ::= `vector.scalable.extract` $source `[` $pos + // `]` attr-dict + // `:` type($res) `from` type($source) + seq(choice('vector.extract', 'vector.load', + 'vector.scalable.extract'), + field('operand', $.value_use), + field('indices', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.fma` $lhs `,` $rhs `,` $acc + // attr-dict `:` type($lhs) + seq('vector.fma', field('lhs', $.value_use), ',', + field('rhs', $.value_use), ',', + field('acc', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.flat_transpose` $matrix attr-dict + // `:` type($matrix) `->` type($res) + seq('vector.flat_transpose', field('matrix', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `vector.insert` $source `,` $dest + // $position attr-dict + // `:` type($source) `into` type($dest) + // operation ::= `vector.scalable.insert` $source `,` $dest + // `[` $pos `]` attr-dict + // `:` type($source) `into` type($dest) + // operation ::= `vector.shuffle` operands $mask attr-dict + // `:` type(operands) operation ::= `vector.store` + // $valueToStore `,` $base `[` $indices `]` attr-dict + // `:` type($base) `,` type($valueToStore) + seq(choice('vector.insert', 'vector.scalable.insert', + 'vector.shuffle', 'vector.store'), + field('source', $.value_use), ',', + field('destination', $.value_use), + field('position', $._dense_idx_list), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.insert_strided_slice` $source `,` + // $dest attr-dict + // `:` type($source) `into` type($dest) + seq('vector.insert_strided_slice', + field('source', $.value_use), ',', + field('destination', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.matrix_multiply` $lhs `,` $rhs + // attr-dict + // `:` `(` type($lhs) `,` type($rhs) `)` + // `->` type($res) + seq('vector.matrix_multiply', field('lhs', $.value_use), + ',', field('rhs', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._function_type_annotation)), + + // operation ::= `vector.print` $source attr-dict `:` + // type($source) + seq(choice('vector.print', 'vector.splat'), + field('operand', $.value_use), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + seq('vector.transfer_read', + field('source', seq($.value_use, $._dense_idx_list)), + field('paddingMask', repeat(seq(',', $.value_use))), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + seq('vector.transfer_write', field('vector', $.value_use), + ',', + field('source', seq($.value_use, $._dense_idx_list)), + field('mask', optional(seq(',', $.value_use))), + field('attributes', optional($.attribute)), + field('return', $._type_annotation)), + + // operation ::= `vector.yield` attr-dict ($operands^ `:` + // type($operands))? + seq('vector.yield', + field('attributes', optional($.attribute)), + field('results', optional($._value_use_type_list))))) +} diff --git a/mlir/utils/tree-sitter-mlir/grammar.js b/mlir/utils/tree-sitter-mlir/grammar.js --- a/mlir/utils/tree-sitter-mlir/grammar.js +++ b/mlir/utils/tree-sitter-mlir/grammar.js @@ -1,261 +1,470 @@ -module.exports = grammar({ +'use strict'; + +const builtin_dialect = require('./dialect/builtin'); +const func_dialect = require('./dialect/func'); +const llvm_dialect = require('./dialect/llvm'); +const arith_dialect = require('./dialect/arith'); +const math_dialect = require('./dialect/math'); +const cf_dialect = require('./dialect/cf'); +const scf_dialect = require('./dialect/scf'); +const memref_dialect = require('./dialect/memref'); +const vector_dialect = require('./dialect/vector'); +const tensor_dialect = require('./dialect/tensor'); +const bufferization_dialect = require('./dialect/bufferization'); +const affine_dialect = require('./dialect/affine'); +const linalg_dialect = require('./dialect/linalg'); + +const common = { + // Top level production: + // (operation | attribute-alias-def | type-alias-def) + toplevel : $ => seq($._toplevel, repeat($._toplevel)), + _toplevel : $ => choice($.operation, $.attribute_alias_def, $.type_alias_def), + + // Common syntax (lang-ref) + // digit ::= [0-9] + // hex_digit ::= [0-9a-fA-F] + // letter ::= [a-zA-Z] + // id-punct ::= [$._-] + // + // integer-literal ::= decimal-literal | hexadecimal-literal + // decimal-literal ::= digit+ + // hexadecimal-literal ::= `0x` hex_digit+ + // float-literal ::= [-+]?[0-9]+[.][0-9]*([eE][-+]?[0-9]+)? + // string-literal ::= `"` [^"\n\f\v\r]* `"` TODO: define escaping rules + // + _digit : $ => /[0-9]/, + integer_literal : $ => choice($._decimal_literal, $._hexadecimal_literal), + _decimal_literal : $ => token(seq(optional(/[-+]/), repeat1(/[0-9]/))), + _hexadecimal_literal : $ => token(seq('0x', repeat1(/[0-9a-fA-F]/))), + float_literal : $ => + token(seq(optional(/[-+]/), repeat1(/[0-9]/), '.', repeat(/[0-9]/), + optional(seq(/[eE]/, optional(/[-+]/), repeat1(/[0-9]/))))), + string_literal : $ => token(seq('"', repeat(/[^\\"\n\f\v\r]+/), '"')), + bool_literal : $ => token(choice('true', 'false')), + unit_literal : $ => token('unit'), + complex_literal : $ => + seq('(', choice($.integer_literal, $.float_literal), ',', + choice($.integer_literal, $.float_literal), ')'), + tensor_literal : $ => + seq(token(choice('dense', 'sparse')), '<', + optional(choice( + seq($.nested_idx_list, repeat(seq(',', $.nested_idx_list))), + $._primitive_idx_literal)), + '>'), + array_literal : $ => seq(token('array'), '<', $.type, ':', $._idx_list, '>'), + _literal : $ => choice($.integer_literal, $.float_literal, $.string_literal, + $.bool_literal, $.tensor_literal, $.array_literal, + $.complex_literal, $.unit_literal), + + nested_idx_list : $ => + seq('[', optional(choice($.nested_idx_list, $._idx_list)), + repeat(seq(',', $.nested_idx_list)), ']'), + _idx_list : $ => prec.right(seq($._primitive_idx_literal, + repeat(seq(',', $._primitive_idx_literal)))), + _primitive_idx_literal : $ => choice($.integer_literal, $.float_literal, + $.bool_literal, $.complex_literal), + + // Identifiers + // bare-id ::= (letter|[_]) (letter|digit|[_$.])* + // bare-id-list ::= bare-id (`,` bare-id)* + // value-id ::= `%` suffix-id + // suffix-id ::= (digit+ | ((letter|id-punct) (letter|id-punct|digit)*)) + // alias-name :: = bare-id + // + // symbol-ref-id ::= `@` (suffix-id | string-literal) (`::` + // symbol-ref-id)? + // value-id-list ::= value-id (`,` value-id)* + // + // // Uses of value, e.g. in an operand list to an operation. + // value-use ::= value-id + // value-use-list ::= value-use (`,` value-use)* + bare_id : $ => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.]/))), + _alias_or_dialect_id : $ => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/))), + bare_id_list : $ => seq($.bare_id, repeat(seq(',', $.bare_id))), + value_use : $ => seq('%', $._suffix_id), + _suffix_id : $ => token(seq( + choice(repeat1(/[0-9]/), seq(/[a-zA-Z_$.-]/, repeat(/[a-zA-Z0-9_$.-]/))), + optional(seq(choice(':', '#'), repeat1(/[0-9]/))))), + symbol_ref_id : $ => seq('@', choice($._suffix_id, $.string_literal), + optional(seq('::', $.symbol_ref_id))), + _value_use_list : $ => seq($.value_use, repeat(seq(',', $.value_use))), + + // Operations + // operation ::= op-result-list? (generic-operation | + // custom-operation) + // trailing-location? + // generic-operation ::= string-literal `(` value-use-list? `)` + // successor-list? region-list? + // dictionary-attribute? `:` function-type + // custom-operation ::= bare-id custom-operation-format + // op-result-list ::= op-result (`,` op-result)* `=` + // op-result ::= value-id (`:` integer-literal) + // successor-list ::= `[` successor (`,` successor)* `]` + // successor ::= caret-id (`:` bb-arg-list)? + // region-list ::= `(` region (`,` region)* `)` + // dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? + // `}` + // trailing-location ::= (`loc` `(` location `)`)? + operation : $ => + seq(field('lhs', optional($._op_result_list)), + field('rhs', choice($.generic_operation, $.custom_operation)), + field('location', optional($.trailing_location))), + generic_operation : $ => seq( + $.string_literal, $._value_use_list_parens, optional($._successor_list), + optional($._region_list), optional($.attribute), ':', $.function_type), + // custom-operation rule is defined later in the grammar, post the generic. + _op_result_list : $ => seq($.op_result, repeat(seq(',', $.op_result)), '='), + op_result : $ => seq($.value_use, optional(seq(':', $.integer_literal))), + _successor_list : $ => + seq('[', $.successor, repeat(seq(',', $.successor)), ']'), + successor : $ => seq($.caret_id, optional($._value_arg_list)), + _region_list : $ => seq('(', $.region, repeat(seq(',', $.region)), ')'), + dictionary_attribute : $ => seq('{', optional($.attribute_entry), + repeat(seq(',', $.attribute_entry)), '}'), + trailing_location : $ => seq(token('loc'), '(', $.location, ')'), + // TODO: Complete location forms. + location : $ => $.string_literal, + + // Blocks + // block ::= block-label operation+ + // block-label ::= block-id block-arg-list? `:` + // block-id ::= caret-id + // caret-id ::= `^` suffix-id + // value-id-and-type ::= value-id `:` type + // + // // Non-empty list of names and types. + // value-id-and-type-list ::= value-id-and-type (`,` value-id-and-type)* + // + // block-arg-list ::= `(` value-id-and-type-list? `)` + block : $ => seq($.block_label, repeat1($.operation)), + block_label : $ => seq($._block_id, optional($.block_arg_list), ':'), + _block_id : $ => $.caret_id, + caret_id : $ => seq('^', $._suffix_id), + _value_use_and_type : $ => seq($.value_use, optional(seq(':', $.type))), + _value_use_and_type_list : $ => + seq($._value_use_and_type, repeat(seq(',', $._value_use_and_type))), + block_arg_list : $ => seq('(', optional($._value_use_and_type_list), ')'), + _value_arg_list : $ => seq('(', optional($._value_use_type_list), ')'), + _value_use_type_list : $ => seq($._value_use_list, $._type_annotation), + + // Regions + // region ::= `{` entry-block? block* `}` + // entry-block ::= operation+ + region : $ => seq('{', optional($.entry_block), repeat($.block), '}'), + entry_block : $ => repeat1($.operation), + + // Types + // type ::= type-alias | dialect-type | builtin-type + // + // type-list-no-parens ::= type (`,` type)* + // type-list-parens ::= `(` type-list-no-parens? `)` + // + // // This is a common way to refer to a value with a specified type. + // ssa-use-and-type ::= ssa-use `:` type + // ssa-use ::= value-use + // + // // Non-empty list of names and types. + // ssa-use-and-type-list ::= ssa-use-and-type (`,` ssa-use-and-type)* + // + // function-type ::= (type | type-list-parens) `->` (type | + // type-list-parens) + type : $ => choice($.type_alias, $.dialect_type, $.builtin_type), + _type_list_no_parens : $ => prec.left(seq($.type, repeat(seq(',', $.type)))), + _type_list_parens : $ => seq('(', optional($._type_list_no_parens), ')'), + function_type : $ => + seq(choice($.type, $._type_list_parens), $._function_return), + _function_return : $ => seq(token('->'), choice($.type, $._type_list_parens)), + _type_annotation : $ => + seq(':', choice(seq($.type, choice('from', 'into', 'to'), $.type), + $._type_list_no_parens)), + _function_type_annotation : $ => seq(':', $.function_type), + _literal_and_type : $ => seq($._literal, optional($._type_annotation)), + + // Type aliases + // type-alias-def ::= '!' alias-name '=' type + // type-alias ::= '!' alias-name + type_alias_def : $ => seq('!', $._alias_or_dialect_id, '=', $.type), + type_alias : $ => seq('!', $._alias_or_dialect_id), + + // Dialect Types + // dialect-namespace ::= bare-id + // + // opaque-dialect-item ::= dialect-namespace '<' string-literal '>' + // + // pretty-dialect-item ::= dialect-namespace '.' + // pretty-dialect-item-lead-ident pretty-dialect-item-body? + // + // pretty-dialect-item-lead-ident ::= '[A-Za-z][A-Za-z0-9._]*' + // pretty-dialect-item-body ::= '<' pretty-dialect-item-contents+ '>' + // pretty-dialect-item-contents ::= pretty-dialect-item-body + // | '(' pretty-dialect-item-contents+ ')' + // | '[' pretty-dialect-item-contents+ ']' + // | '{' pretty-dialect-item-contents+ '}' + // | '[^[<({>\])}\0]+' + // + // dialect-type ::= '!' (opaque-dialect-item | pretty-dialect-item) + dialect_type : $ => + seq('!', choice($.opaque_dialect_item, $.pretty_dialect_item)), + dialect_namespace : $ => $._alias_or_dialect_id, + dialect_ident : $ => $._alias_or_dialect_id, + opaque_dialect_item : $ => + seq($.dialect_namespace, '<', $.string_literal, '>'), + pretty_dialect_item : $ => seq($.dialect_namespace, '.', $.dialect_ident, + optional($.pretty_dialect_item_body)), + pretty_dialect_item_body : $ => + seq('<', repeat($._pretty_dialect_item_contents), '>'), + _pretty_dialect_item_contents : $ => + prec.left(choice($.pretty_dialect_item_body, repeat1(/[^<>]/))), + + // Builtin types + builtin_type : $ => choice( + // TODO: Add opaque_type + $.integer_type, $.float_type, $.complex_type, $.index_type, $.memref_type, + $.none_type, $.tensor_type, $.vector_type, $.tuple_type), + + // signed-integer-type ::= `si`[1-9][0-9]* + // unsigned-integer-type ::= `ui`[1-9][0-9]* + // signless-integer-type ::= `i`[1-9][0-9]* + // integer-type ::= signed-integer-type | unsigned-integer-type | + // signless-integer-type + integer_type : $ => + token(seq(choice('si', 'ui', 'i'), /[1-9]/, repeat(/[0-9]/))), + float_type : $ => token( + choice('f16', 'f32', 'f64', 'f80', 'f128', 'bf16', 'f8E4M3FN', 'f8E5M2')), + index_type : $ => token('index'), + none_type : $ => token('none'), + complex_type : $ => seq(token('complex'), '<', $._prim_type, '>'), + _prim_type : $ => + choice($.integer_type, $.float_type, $.index_type, $.complex_type, + $.none_type, $.memref_type, $.vector_type, $.tensor_type), + + // memref-type ::= `memref` `<` dimension-list-ranked type + // (`,` layout-specification)? (`,` memory-space)? `>` + // layout-specification ::= attribute-value + // memory-space ::= attribute-value + memref_type : $ => + seq(token('memref'), '<', field('dimension_list', $.dim_list), + optional(seq(',', $.attribute_value)), + optional(seq(',', $.attribute_value)), '>'), + dim_list : $ => seq($._dim_primitive, repeat(seq('x', $._dim_primitive))), + _dim_primitive : $ => choice($._prim_type, repeat1($._digit), '?', '*'), + + // tensor-type ::= `tensor` `<` dimension-list type (`,` encoding)? `>` + // dimension-list ::= (dimension `x`)* + // dimension ::= `?` | decimal-literal + // encoding ::= attribute-value + // tensor-type ::= `tensor` `<` `*` `x` type `>` + tensor_type : $ => seq(token('tensor'), '<', $.dim_list, + optional(seq(',', $.tensor_encoding)), '>'), + tensor_encoding : $ => $.attribute_value, + + // vector-type ::= `vector` `<` vector-dim-list vector-element-type `>` + // vector-element-type ::= float-type | integer-type | index-type + // vector-dim-list := (static-dim-list `x`)? (`[` static-dim-list `]` `x`)? + // static-dim-list ::= decimal-literal (`x` decimal-literal)* + vector_type : $ => + seq(token('vector'), '<', optional($.vector_dim_list), $._prim_type, '>'), + vector_dim_list : $ => + choice(seq($._static_dim_list, 'x', + optional(seq('[', $._static_dim_list, ']', 'x'))), + seq('[', $._static_dim_list, ']', 'x')), + _static_dim_list : $ => + seq(repeat1($._digit), repeat(seq('x', repeat1($._digit)))), + + // tuple-type ::= `tuple` `<` (type ( `,` type)*)? `>` + tuple_type : $ => + seq(token('tuple'), '<', $.tuple_dim, repeat(seq(',', $.tuple_dim)), '>'), + tuple_dim : $ => $._prim_type, + + // Attributes + // attribute-entry ::= (bare-id | string-literal) `=` attribute-value + // attribute-value ::= attribute-alias | dialect-attribute | + // builtin-attribute + attribute_entry : $ => seq(choice($.bare_id, $.string_literal), + optional(seq('=', $.attribute_value))), + attribute_value : $ => + choice(seq('[', optional($._attribute_value_nobracket), + repeat(seq(',', $._attribute_value_nobracket)), ']'), + $._attribute_value_nobracket), + _attribute_value_nobracket : $ => + choice($.attribute_alias, $.dialect_attribute, $.builtin_attribute, + $.dictionary_attribute, $._literal_and_type, $.type), + attribute : $ => choice($.attribute_alias, $.dialect_attribute, + $.builtin_attribute, $.dictionary_attribute), + + // Attribute Value Aliases + // attribute-alias-def ::= '#' alias-name '=' attribute-value + // attribute-alias ::= '#' alias-name + attribute_alias_def : $ => + seq('#', $._alias_or_dialect_id, '=', $.attribute_value), + attribute_alias : $ => seq('#', $._alias_or_dialect_id), + + // Dialect Attribute Values + dialect_attribute : $ => + seq('#', choice($.opaque_dialect_item, $.pretty_dialect_item)), + + // Builtin Attribute Values + builtin_attribute : $ => choice( + // TODO + $.strided_layout, $.affine_map, $.affine_set), + strided_layout : $ => seq(token('strided'), '<', '[', $._dim_list_comma, ']', + optional(seq(',', token('offset'), ':', + choice($.integer_literal, '?', '*'))), + '>'), + _dim_list_comma : $ => + seq($._dim_primitive, repeat(seq(',', $._dim_primitive))), + + affine_map : $ => + seq(token('affine_map'), '<', $._multi_dim_affine_expr_parens, + optional($._multi_dim_affine_expr_sq), token('->'), + $._multi_dim_affine_expr_parens, '>'), + affine_set : $ => + seq(token('affine_set'), '<', $._multi_dim_affine_expr_parens, + optional($._multi_dim_affine_expr_sq), ':', + $._multi_dim_affine_expr_parens, '>'), + _multi_dim_affine_expr_parens : $ => + seq('(', optional($._multi_dim_affine_expr), ')'), + _multi_dim_affine_expr_sq : $ => + seq('[', optional($._multi_dim_affine_expr), ']'), + + // affine-expr ::= `(` affine-expr `)` + // | affine-expr `+` affine-expr + // | affine-expr `-` affine-expr + // | `-`? integer-literal `*` affine-expr + // | affine-expr `ceildiv` integer-literal + // | affine-expr `floordiv` integer-literal + // | affine-expr `mod` integer-literal + // | `-`affine-expr + // | bare-id + // | `-`? integer-literal + // multi-dim-affine-expr ::= `(` `)` + // | `(` affine-expr (`,` affine-expr)* `)` + + // semi-affine-expr ::= `(` semi-affine-expr `)` + // | semi-affine-expr `+` semi-affine-expr + // | semi-affine-expr `-` semi-affine-expr + // | symbol-or-const `*` semi-affine-expr + // | semi-affine-expr `ceildiv` symbol-or-const + // | semi-affine-expr `floordiv` symbol-or-const + // | semi-affine-expr `mod` symbol-or-const + // | bare-id + // | `-`? integer-literal + // symbol-or-const ::= `-`? integer-literal | symbol-id + // multi-dim-semi-affine-expr ::= `(` semi-affine-expr (`,` semi-affine-expr)* + // `)` + + // affine-constraint ::= affine-expr `>=` `affine-expr` + // | affine-expr `<=` `affine-expr` + // | affine-expr `==` `affine-expr` + // affine-constraint-conjunction ::= affine-constraint (`,` + // affine-constraint)* + + _multi_dim_affine_expr : $ => + seq($._affine_expr, repeat(seq(',', $._affine_expr))), + _affine_expr : $ => prec.right(choice( + seq('(', $._affine_expr, ')'), seq('-', $._affine_expr), + seq($._affine_expr, $._affine_token, $._affine_expr), $._affine_prim)), + _affine_prim : $ => + choice($.integer_literal, $.value_use, $.bare_id, + seq('symbol', '(', $.value_use, ')'), + seq(choice('max', 'min'), '(', $._value_use_list, ')')), + _affine_token : $ => token( + choice('+', '-', '*', 'ceildiv', 'floordiv', 'mod', '==', '>=', '<=')), + + func_return : $ => seq(token('->'), $.type_list_attr_parens), + func_arg_list : $ => seq( + '(', optional(choice($.variadic, $._value_id_and_type_attr_list)), ')'), + _value_id_and_type_attr_list : $ => seq( + $._value_id_and_type_attr, repeat(seq(',', $._value_id_and_type_attr)), + optional(seq(',', $.variadic))), + _value_id_and_type_attr : $ => seq($._function_arg, optional($.attribute)), + _function_arg : $ => + choice(seq($.value_use, ':', $.type), $.value_use, $.type), + type_list_attr_parens : $ => + choice($.type, + seq('(', $.type, optional($.attribute), + repeat(seq(',', $.type, optional($.attribute))), ')'), + seq('(', ')')), + variadic : $ => token('...'), + + // (func.func|llvm.func) takes arguments, an optional return type, and and + // optional body + _op_func : $ => + seq(field('visibility', optional('private')), + field('name', $.symbol_ref_id), field('arguments', $.func_arg_list), + field('return', optional($.func_return)), + field('attributes', optional(seq(token('attributes'), $.attribute))), + field('body', optional($.region))), + + // dim-use-list ::= `(` ssa-use-list? `)` + // symbol-use-list ::= `[` ssa-use-list? `]` + // dim-and-symbol-use-list ::= dim-use-list symbol-use-list? + _value_use_list_parens : $ => seq('(', optional($._value_use_list), ')'), + _dim_and_symbol_use_list : $ => + seq($._value_use_list_parens, optional($._dense_idx_list)), + + // assignment-list ::= assignment | assignment `,` assignment-list + // assignment ::= ssa-value `=` ssa-value + _value_assignment_list : $ => seq('(', optional($._value_assignment), + repeat(seq(',', $._value_assignment)), ')'), + _value_assignment : $ => seq($.value_use, '=', $.value_use), + + _dense_idx_list : $ => seq( + '[', + optional(seq(choice($.integer_literal, $.value_use), + repeat(seq(',', choice($.integer_literal, $.value_use))))), + ']'), + + // lower-bound ::= `max`? affine-map-attribute dim-and-symbol-use-list | + // shorthand-bound + // upper-bound ::= `min`? affine-map-attribute dim-and-symbol-use-list | + // shorthand-bound + // shorthand-bound ::= ssa-id | `-`? integer-literal + _bound : $ => + choice(seq($.attribute, $._dim_and_symbol_use_list), $._shorthand_bound), + _shorthand_bound : $ => choice($.value_use, $.integer_literal), + + // Dialect-specific attributes + restrict_attr : $ => token('restrict'), + writable_attr : $ => token('writable'), + gather_dims_attr : $ => + seq(token('gather_dims'), '(', $._dense_idx_list, ')'), + scatter_dims_attr : $ => + seq(token('scatter_dims'), '(', $._dense_idx_list, ')'), + unique_attr : $ => token('unique'), + nofold_attr : $ => token('nofold'), + outer_dims_perm_attr : $ => + seq(token('outer_dims_perm'), '=', $._dense_idx_list), + inner_dims_pos_attr : $ => + seq(token('inner_dims_pos'), '=', $._dense_idx_list), + inner_tiles_attr : $ => seq(token('inner_tiles'), '=', $._dense_idx_list), + isWrite_attr : $ => token(choice('read', 'write')), + localityHint_attr : $ => seq(token('locality'), '<', $.integer_literal, '>'), + isDataCache_attr : $ => token(choice('data', 'instr')), + fastmath_attr : $ => + seq(token('fastmath'), '<', + seq($._fastmath_flag, repeat(seq(',', $._fastmath_flag))), '>'), + _fastmath_flag : $ => token(choice('none', 'reassoc', 'nnan', 'ninf', 'nsz', + 'arcp', 'contract', 'afn', 'fast')), + + // Comment (standard BCPL) + comment : $ => token(seq('//', /.*/)), + + // TODO: complete + custom_operation : $ => + choice($.builtin_dialect, $.func_dialect, $.llvm_dialect, $.arith_dialect, + $.math_dialect, $.cf_dialect, $.scf_dialect, $.memref_dialect, + $.vector_dialect, $.tensor_dialect, $.bufferization_dialect, + $.affine_dialect, $.linalg_dialect) +} + + module.exports = grammar({ name : 'mlir', - extras : $ => [/\s/, - $.comment, -], - conflicts : $ => [], - rules : { - // Top level production: - // (operation | attribute-alias-def | type-alias-def) - toplevel : $ => seq(choice( - $.operation, - $.attribute_alias_def, - $.type_alias_def, - )), - - // Common syntax (lang-ref) - // digit ::= [0-9] - // hex_digit ::= [0-9a-fA-F] - // letter ::= [a-zA-Z] - // id-punct ::= [$._-] - // - // integer-literal ::= decimal-literal | hexadecimal-literal - // decimal-literal ::= digit+ - // hexadecimal-literal ::= `0x` hex_digit+ - // float-literal ::= [-+]?[0-9]+[.][0-9]*([eE][-+]?[0-9]+)? - // string-literal ::= `"` [^"\n\f\v\r]* `"` TODO: define escaping rules - // - _digit : $ => /[0-9]/, - _hex_digit : $ => /[0-9a-fA-F]/, - integer_literal : $ => choice($._decimal_literal, $._hexadecimal_literal), - _decimal_literal : $ => repeat1($._digit), - _hexadecimal_literal : $ => seq('0x', repeat1($._hex_digit)), - float_literal : $ => token( - seq(optional(/[-+]/), repeat1(/[0_9]/), - optional(seq('.', repeat(/[0-9]/), - optional(seq(/[eE]/, optional(/[-+]/), - repeat1(/[0-9]/))))))), - string_literal : $ => seq( - '"', - repeat(token.immediate(prec(1, /[^\\"\n\f\v\r]+/))), - '"', - ), - - // Identifiers - // bare-id ::= (letter|[_]) (letter|digit|[_$.])* - // bare-id-list ::= bare-id (`,` bare-id)* - // value-id ::= `%` suffix-id - // suffix-id ::= (digit+ | ((letter|id-punct) (letter|id-punct|digit)*)) - // alias-name :: = bare-id - // - // symbol-ref-id ::= `@` (suffix-id | string-literal) (`::` - // symbol-ref-id)? - // value-id-list ::= value-id (`,` value-id)* - // - // // Uses of value, e.g. in an operand list to an operation. - // value-use ::= value-id - // value-use-list ::= value-use (`,` value-use)* - bare_id : $ => seq(token(/[a-zA-Z_]/), - token.immediate(repeat(/[a-zA-Z0-9_$]/))), - bare_id_list : $ => seq($.bare_id, repeat(seq(',', $.bare_id))), - value_id : $ => seq('%', $._suffix_id), - alias_name : $ => $.bare_id, - _suffix_id : $ => choice(repeat1(/[0-9]/), - seq(/[a-zA-Z_$.]/, repeat(/[a-zA-Z0-9_$.]/))), - symbol_ref_id : $ => seq('@', choice($._suffix_id, $.string_literal), - optional(seq('::', $.symbol_ref_id))), - value_use : $ => $.value_id, - value_use_list : $ => seq($.value_use, repeat(seq(',', $.value_use))), - - // Operations - // operation ::= op-result-list? (generic-operation | - // custom-operation) - // trailing-location? - // generic-operation ::= string-literal `(` value-use-list? `)` - // successor-list? - // region-list? dictionary-attribute? `:` - // function-type - // custom-operation ::= bare-id custom-operation-format - // op-result-list ::= op-result (`,` op-result)* `=` - // op-result ::= value-id (`:` integer-literal) - // successor-list ::= `[` successor (`,` successor)* `]` - // successor ::= caret-id (`:` bb-arg-list)? - // region-list ::= `(` region (`,` region)* `)` - // dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? - // `}` - // trailing-location ::= (`loc` `(` location `)`)? - operation : $ => seq(optional($.op_result_list), - choice($.generic_operation, $.custom_operation), - optional($.trailing_location)), - generic_operation : $ => - seq($.string_literal, '(', optional($.value_use_list), - ')', optional($.successor_list), - optional($.region_list), - optional($.dictionary_attribute), ':', - $.function_type), - // custom-operation rule is defined later in the grammar, post the generic. - op_result_list : $ => seq($.op_result, repeat(seq(',', $.op_result)), '='), - op_result : $ => seq($.value_id, optional(seq(':', $.integer_literal))), - successor_list : $ => seq('[', $.successor, repeat(seq(',', $.successor)), - ']'), - successor : $ => seq($.caret_id, optional(seq(':', $.block_arg_list))), - region_list : $ => seq('(', $.region, repeat(seq(',', $.region)), ')'), - dictionary_attribute : $ => seq( - '{', - optional(seq($.attribute_entry, - repeat(seq(',', $.attribute_entry)))), - '}'), - trailing_location : $ => seq('loc(', $.location, ')'), - // TODO: Complete location forms. - location : $ => $.string_literal, - - // Blocks - // block ::= block-label operation+ - // block-label ::= block-id block-arg-list? `:` - // block-id ::= caret-id - // caret-id ::= `^` suffix-id - // value-id-and-type ::= value-id `:` type - // - // // Non-empty list of names and types. - // value-id-and-type-list ::= value-id-and-type (`,` value-id-and-type)* - // - // block-arg-list ::= `(` value-id-and-type-list? `)` - block : $ => seq($.block_label, repeat1($.operation)), - block_label : $ => seq($._block_id, optional($.block_arg_list), ':'), - _block_id : $ => $.caret_id, - caret_id : $ => seq('^', $._suffix_id), - value_id_and_type : $ => seq($.value_id, ':', $.type), - value_id_and_type_list : $ => seq($.value_id_and_type, - repeat(seq(',', $.value_id_and_type))), - block_arg_list : $ => seq('(', optional($.value_id_and_type_list), ')'), - - // Regions - // region ::= `{` entry-block? block* `}` - // entry-block ::= operation+ - region : $ => seq('{', optional($.entry_block), repeat($.block), '}'), - entry_block : $ => repeat1($.operation), - - // Types - // type ::= type-alias | dialect-type | builtin-type - // - // type-list-no-parens ::= type (`,` type)* - // type-list-parens ::= `(` type-list-no-parens? `)` - // - // // This is a common way to refer to a value with a specified type. - // ssa-use-and-type ::= ssa-use `:` type - // ssa-use ::= value-use - // - // // Non-empty list of names and types. - // ssa-use-and-type-list ::= ssa-use-and-type (`,` ssa-use-and-type)* - // - // function-type ::= (type | type-list-parens) `->` (type | - // type-list-parens) - type : $ => choice($.type_alias, $.dialect_type, $.builtin_type), - type_list_no_parens : $ => seq($.type, repeat(seq(',', $.type))), - type_list_parens : $ => seq('(', optional($.type_list_no_parens), ')'), - ssa_use_and_type : $ => seq($.ssa_use, ':', $.type), - ssa_use : $ => $.value_use, - ssa_use_and_type_list : $ => seq($.ssa_use_and_type, - repeat(seq(',', $.ssa_use_and_type))), - function_type : $ => seq(choice($.type, $.type_list_parens), '->', - choice($.type, $.type_list_parens)), - - // Type aliases - // type-alias-def ::= '!' alias-name '=' type - // type-alias ::= '!' alias-name - type_alias_def : $ => seq('!', $.alias_name, '=', $.type), - type_alias : $ => seq('!', $.alias_name), - - // Dialect Types - // dialect-namespace ::= bare-id - // - // opaque-dialect-item ::= dialect-namespace '<' string-literal '>' - // - // pretty-dialect-item ::= dialect-namespace '.' - // pretty-dialect-item-lead-ident - // pretty-dialect-item-body? - // - // pretty-dialect-item-lead-ident ::= '[A-Za-z][A-Za-z0-9._]*' - // pretty-dialect-item-body ::= '<' pretty-dialect-item-contents+ '>' - // pretty-dialect-item-contents ::= pretty-dialect-item-body - // | '(' pretty-dialect-item-contents+ ')' - // | '[' pretty-dialect-item-contents+ ']' - // | '{' pretty-dialect-item-contents+ '}' - // | '[^[<({>\])}\0]+' - // - // dialect-type ::= '!' (opaque-dialect-item | pretty-dialect-item) - dialect_type : $ => seq( - '!', choice($.opaque_dialect_item, $.pretty_dialect_item)), - dialect_namespace : $ => $.bare_id, - opaque_dialect_item : $ => seq($.dialect_namespace, '<', $.string_literal, - '>'), - pretty_dialect_item : $ => seq($.dialect_namespace, '.', - $.pretty_dialect_item_lead_ident, - optional($.pretty_dialect_item_body)), - pretty_dialect_item_lead_ident : $ => $.bare_id, - pretty_dialect_item_body : $ => seq('<', - repeat1($.pretty_dialect_item_contents), - '>'), - // TODO: not sure why prec.left (setting left-associated parsing) needed - // here, - // left-associated way avoids an ambiguity flagged by generator. It may not - // be needed and be only papering over an issue. - pretty_dialect_item_contents : $ => prec.left(choice( - $.pretty_dialect_item_body, - seq('(', - repeat1( - $.pretty_dialect_item_contents), - ')'), - seq('[', - repeat1( - $.pretty_dialect_item_contents), - ']'), - seq('{', - repeat1( - $.pretty_dialect_item_contents), - '}'), - repeat1(/[^\[<({>\])}\\0]/))), - dialect_type : $ => seq( - '!', choice($.opaque_dialect_item, $.pretty_dialect_item)), - - // Builtin types - builtin_type : $ => choice( - // TODO: Add builtin types - seq('i', repeat1(/[0-9]/))), - - // Attributes - // attribute-entry ::= (bare-id | string-literal) `=` attribute-value - // attribute-value ::= attribute-alias | dialect-attribute | - // builtin-attribute - attribute_entry : $ => seq(choice($.bare_id, $.string_literal), '=', - $.attribute_value), - attribute_value : $ => choice($.attribute_alias, $.dialect_attribute, - $.builtin_attribute), - - // Attribute Value Aliases - // attribute-alias-def ::= '#' alias-name '=' attribute-value - // attribute-alias ::= '#' alias-name - attribute_alias_def : $ => seq('#', $.alias_name, '=', $.attribute_value), - attribute_alias : $ => seq('#', $.alias_name), - - // Dialect Attribute Values - dialect_attribute : $ => seq('#', choice($.opaque_dialect_item, - $.pretty_dialect_item)), - - // Builtin Attribute Values - builtin_attribute : $ => choice( - // TODO - $.function_type, - $.string_literal, - ), - - // Comment (standard BCPL) - comment : $ => token(seq('//', /.*/)), - - custom_operation : $ => choice( - // TODO: Just basic/incomplete instance. - seq('func', field('name', $.symbol_ref_id), - $.block_arg_list, '->', $.type, $.region), - ), - } + extras : $ => [/\s/, $.comment], + conflicts : $ => [[ $._static_dim_list, $._static_dim_list ], + [ $.dictionary_attribute, $.region ]], + rules : Object.assign(common, builtin_dialect, func_dialect, llvm_dialect, + arith_dialect, math_dialect, cf_dialect, scf_dialect, + memref_dialect, vector_dialect, tensor_dialect, + bufferization_dialect, affine_dialect, linalg_dialect) }); diff --git a/mlir/utils/tree-sitter-mlir/package-lock.json b/mlir/utils/tree-sitter-mlir/package-lock.json new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/package-lock.json @@ -0,0 +1,118 @@ +{ + "name": "tree-sitter-mlir", + "version": "0.0.1", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "tree-sitter-mlir", + "version": "0.0.1", + "license": "Apache-2.0", + "dependencies": { + "nan": "^2.17.0" + }, + "devDependencies": { + "glob": "^8.1.0", + "tree-sitter-cli": "^0.20.7" + } + }, + "node_modules/balanced-match": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", + "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==", + "dev": true + }, + "node_modules/brace-expansion": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.1.tgz", + "integrity": "sha512-XnAIvQ8eM+kC6aULx6wuQiwVsnzsi9d3WxzV3FpWTGA19F621kwdbsAcFKXgKUHZWsy+mY6iL1sHTxWEFCytDA==", + "dev": true, + "dependencies": { + "balanced-match": "^1.0.0" + } + }, + "node_modules/fs.realpath": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/fs.realpath/-/fs.realpath-1.0.0.tgz", + "integrity": "sha512-OO0pH2lK6a0hZnAdau5ItzHPI6pUlvI7jMVnxUQRtw4owF2wk8lOSabtGDCTP4Ggrg2MbGnWO9X8K1t4+fGMDw==", + "dev": true + }, + "node_modules/glob": { + "version": "8.1.0", + "resolved": "https://registry.npmjs.org/glob/-/glob-8.1.0.tgz", + "integrity": "sha512-r8hpEjiQEYlF2QU0df3dS+nxxSIreXQS1qRhMJM0Q5NDdR386C7jb7Hwwod8Fgiuex+k0GFjgft18yvxm5XoCQ==", + "dev": true, + "dependencies": { + "fs.realpath": "^1.0.0", + "inflight": "^1.0.4", + "inherits": "2", + "minimatch": "^5.0.1", + "once": "^1.3.0" + }, + "engines": { + "node": ">=12" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" + } + }, + "node_modules/inflight": { + "version": "1.0.6", + "resolved": "https://registry.npmjs.org/inflight/-/inflight-1.0.6.tgz", + "integrity": "sha512-k92I/b08q4wvFscXCLvqfsHCrjrF7yiXsQuIVvVE7N82W3+aqpzuUdBbfhWcy/FZR3/4IgflMgKLOsvPDrGCJA==", + "dev": true, + "dependencies": { + "once": "^1.3.0", + "wrappy": "1" + } + }, + "node_modules/inherits": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", + "dev": true + }, + "node_modules/minimatch": { + "version": "5.1.6", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-5.1.6.tgz", + "integrity": "sha512-lKwV/1brpG6mBUFHtb7NUmtABCb2WZZmm2wNiOA5hAb8VdCS4B3dtMWyvcoViccwAW/COERjXLt0zP1zXUN26g==", + "dev": true, + "dependencies": { + "brace-expansion": "^2.0.1" + }, + "engines": { + "node": ">=10" + } + }, + "node_modules/nan": { + "version": "2.17.0", + "resolved": "https://registry.npmjs.org/nan/-/nan-2.17.0.tgz", + "integrity": "sha512-2ZTgtl0nJsO0KQCjEpxcIr5D+Yv90plTitZt9JBfQvVJDS5seMl3FOvsh3+9CoYWXf/1l5OaZzzF6nDm4cagaQ==" + }, + "node_modules/once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "dependencies": { + "wrappy": "1" + } + }, + "node_modules/tree-sitter-cli": { + "version": "0.20.7", + "resolved": "https://registry.npmjs.org/tree-sitter-cli/-/tree-sitter-cli-0.20.7.tgz", + "integrity": "sha512-MHABT8oCPr4D0fatsPo6ATQ9H4h9vHpPRjlxkxJs80tpfAEKGn6A1zU3eqfCKBcgmfZDe9CiL3rKOGMzYHwA3w==", + "dev": true, + "hasInstallScript": true, + "bin": { + "tree-sitter": "cli.js" + } + }, + "node_modules/wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true + } + } +} diff --git a/mlir/utils/tree-sitter-mlir/package.json b/mlir/utils/tree-sitter-mlir/package.json new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/package.json @@ -0,0 +1,34 @@ +{ + "name": "tree-sitter-mlir", + "description": "Tree sitter grammar for MLIR", + "version": "0.0.1", + "license": "Apache-2.0", + "main": "grammar.js", + "scripts": { + "compile": "tree-sitter generate", + "update-tests": "tree-sitter test --update", + "test": "tree-sitter test", + "bench": "node bench.mjs" + }, + "keywords": [ + "tree-sitter", + "lexer", + "parser", + "grammar" + ], + "dependencies": { + "nan": "^2.17.0" + }, + "devDependencies": { + "glob": "^8.1.0", + "tree-sitter-cli": "^0.20.7" + }, + "tree-sitter": [ + { + "scope": "source.mlir", + "file-types": [ + "mlir" + ] + } + ] +} diff --git a/mlir/utils/tree-sitter-mlir/queries/highlights.scm b/mlir/utils/tree-sitter-mlir/queries/highlights.scm new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/queries/highlights.scm @@ -0,0 +1,349 @@ +[ + "ins" + "outs" + "else" + "do" + "loc" + "attributes" + "into" + "to" + "from" + "step" + "low" + "high" + "iter_args" + "padding_value" + "inner_tiles" + "gather_dims" + "scatter_dims" + "outer_dims_perm" + "inner_dims_pos" + "shared_outs" + "default" + + "eq" + "ne" + "oeq" + "olt" + "ole" + "ogt" + "oge" + "slt" + "sle" + "sgt" + "sge" + "ult" + "ule" + "ugt" + "uge" +] @keyword + +[ + "module" + "unrealized_conversion_cast" + + "func.call" + "call" + "func.call_indirect" + "call_indirect" + "func.constant" + "constant" + "func.func" + "func.return" + "return" + + "llvm.func" + "llvm.return" + + "cf.assert" + "cf.br" + "cf.cond_br" + "cf.switch" + + "scf.condition" + "scf.execute_region" + "scf.if" + "scf.index_switch" + "scf.for" + "scf.forall" + "scf.forall.in_parallel" + "scf.parallel" + "scf.reduce" + "scf.reduce.return" + "scf.while" + "scf.yield" + + "arith.constant" + "arith.addi" + "arith.subi" + "arith.divsi" + "arith.divui" + "arith.ceildivsi" + "arith.ceildivui" + "arith.floordivsi" + "arith.remsi" + "arith.remui" + "arith.muli" + "arith.mulsi_extended" + "arith.mului_extended" + "arith.andi" + "arith.ori" + "arith.xori" + "arith.maxsi" + "arith.maxui" + "arith.minsi" + "arith.minui" + "arith.shli" + "arith.shrsi" + "arith.shrui" + "arith.addui_extended" + "arith.addf" + "arith.divf" + "arith.maxf" + "arith.minf" + "arith.mulf" + "arith.remf" + "arith.subf" + "arith.negf" + "arith.cmpi" + "arith.cmpf" + "arith.extf" + "arith.extsi" + "arith.extui" + "arith.fptosi" + "arith.fptoui" + "arith.index_cast" + "arith.index_castui" + "arith.sitofp" + "arith.uitofp" + "arith.bitcast" + "arith.truncf" + "arith.select" + + "math.absf" + "math.atan" + "math.cbrt" + "math.ceil" + "math.cos" + "math.erf" + "math.exp" + "math.exp2" + "math.expm1" + "math.floor" + "math.log" + "math.log10" + "math.log1p" + "math.log2" + "math.round" + "math.roundeven" + "math.rsqrt" + "math.sin" + "math.sqrt" + "math.tan" + "math.tanh" + "math.trunc" + "math.absi" + "math.ctlz" + "math.cttz" + "math.ctpop" + "math.atan2" + "math.copysign" + "math.fpowi" + "math.powf" + "math.ipowi" + "math.fma" + + "memref.alloc" + "memref.cast" + "memref.copy" + "memref.collapse_shape" + "memref.expand_shape" + "memref.prefetch" + "memref.rank" + "memref.realloc" + "memref.view" + + "vector.bitcast" + "vector.broadcast" + "vector.shape_cast" + "vector.type_cast" + "vector.constant_mask" + "vector.create_mask" + "vector.extract" + "vector.load" + "vector.scalable.extract" + "vector.fma" + "vector.flat_transpose" + "vector.insert" + "vector.scalable.insert" + "vector.shuffle" + "vector.store" + "vector.insert_strided_slice" + "vector.matrix_multiply" + "vector.print" + "vector.splat" + "vector.transfer_read" + "vector.transfer_write" + "vector.yield" + + "tensor.empty" + "tensor.cast" + "tensor.dim" + "tensor.collapse_shape" + "tensor.expand_shape" + "tensor.extract" + "tensor.insert" + "tensor.extract_slice" + "tensor.insert_slice" + "tensor.parallel_insert_slice" + "tensor.from_elements" + "tensor.gather" + "tensor.scatter" + "tensor.pad" + "tensor.reshape" + "tensor.splat" + "tensor.pack" + "tensor.unpack" + "tensor.generate" + "tensor.rank" + "tensor.yield" + + "bufferization.alloc_tensor" + "bufferization.to_memref" + "bufferization.to_tensor" + + "linalg.batch_matmul" + "linalg.batch_matmul_transpose_b" + "linalg.batch_matvec" + "linalg.batch_reduce_matmul" + "linalg.broadcast" + "linalg.conv_1d_ncw_fcw" + "linalg.conv_1d_nwc_wcf" + "linalg.conv_1d" + "linalg.conv_2d_nchw_fchw" + "linalg.conv_2d_ngchw_fgchw" + "linalg.conv_2d_nhwc_fhwc" + "linalg.conv_2d_nhwc_hwcf" + "linalg.conv_2d_nhwc_hwcf_q" + "linalg.conv_2d" + "linalg.conv_3d_ndhwc_dhwcf" + "linalg.conv_3d_ndhwc_dhwcf_q" + "linalg.conv_3d" + "linalg.copy" + "linalg.depthwise_conv_1d_nwc_wcm" + "linalg.depthwise_conv_2d_nchw_chw" + "linalg.depthwise_conv_2d_nhwc_hwc" + "linalg.depthwise_conv_2d_nhwc_hwc_q" + "linalg.depthwise_conv_2d_nhwc_hwcm" + "linalg.depthwise_conv_2d_nhwc_hwcm_q" + "linalg.depthwise_conv_3d_ndhwc_dhwc" + "linalg.depthwise_conv_3d_ndhwc_dhwcm" + "linalg.dot" + "linalg.elemwise_binary" + "linalg.elemwise_unary" + "linalg.fill" + "linalg.fill_rng_2d" + "linalg.matmul" + "linalg.matmul_transpose_b" + "linalg.matmul_unsigned" + "linalg.matvec" + "linalg.mmt4d" + "linalg.pooling_nchw_max" + "linalg.pooling_nchw_sum" + "linalg.pooling_ncw_max" + "linalg.pooling_ncw_sum" + "linalg.pooling_ndhwc_max" + "linalg.pooling_ndhwc_min" + "linalg.pooling_ndhwc_sum" + "linalg.pooling_nhwc_max" + "linalg.pooling_nhwc_max_unsigned" + "linalg.pooling_nhwc_min" + "linalg.pooling_nhwc_min_unsigned" + "linalg.pooling_nhwc_sum" + "linalg.pooling_nwc_max" + "linalg.pooling_nwc_max_unsigned" + "linalg.pooling_nwc_min" + "linalg.pooling_nwc_min_unsigned" + "linalg.pooling_nwc_sum" + "linalg.quantized_batch_matmul" + "linalg.quantized_matmul" + "linalg.vecmat" + "linalg.generic" + "linalg.index" + "linalg.map" + "linalg.yield" +] @function.builtin + +(generic_operation) @function + +(builtin_type) @type.builtin + +[ + (type_alias) + (dialect_type) + (type_alias_def) +] @type + +[ + (integer_literal) + (float_literal) + (complex_literal) +] @number + +[ + (bool_literal) + (tensor_literal) + (array_literal) + (unit_literal) +] @constant.builtin + +(string_literal) @string + +[ + (attribute_alias_def) + (attribute_alias) + (bare_attribute_entry) + (attribute) + (fastmath_attr) + (scatter_dims_attr) + (gather_dims_attr) + (outer_dims_perm_attr) + (inner_dims_pos_attr) + (inner_tiles_attr) + (unique_attr) + (nofold_attr) + (isWrite_attr) + (localityHint_attr) + (isDataCache_attr) + (restrict_attr) + (writable_attr) +] @attribute + +[ + "(" + ")" + "{" + "}" + "[" + "]" +] @punctuation.bracket + +[ + ":" + "," +] @punctuation.delimeter + +[ + "=" + "->" +] @operator + +(func_dialect name: (symbol_ref_id) @function) +(llvm_dialect name: (symbol_ref_id) @function) + +(func_arg_list (value_use) @variable.parameter) +(block_arg_list (value_use) @variable.parameter) + +(caret_id) @tag +(value_use) @variable +(comment) @comment diff --git a/mlir/utils/tree-sitter-mlir/queries/locals.scm b/mlir/utils/tree-sitter-mlir/queries/locals.scm new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/queries/locals.scm @@ -0,0 +1,7 @@ +(region) @local.scope + +(func_arg_list (value_use) @local.definition) +(block_arg_list (value_use) @local.definition) +(op_result (value_use) @local.definition) + +(value_use) @local.reference diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/affine.txt b/mlir/utils/tree-sitter-mlir/test/corpus/affine.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/affine.txt @@ -0,0 +1,500 @@ +================================================================================ +affine_map with indexing and sum +================================================================================ +#inline_map_minmax_loop2 = affine_map<()[s0] -> (100, s0 + 1)> +#map0 = affine_map<(i) -> ()> +#bound_map1 = affine_map<(i, j)[s] -> (i + j + s)> +-------------------------------------------------------------------------------- + +(toplevel + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (integer_literal) + (bare_id) + (integer_literal))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id)))))) + +================================================================================ +affine_set expression +================================================================================ +#set0 = affine_set<(i)[N, M] : (i >= 0, -i + N >= 0, N - 5 == 0, -i + M + 1 >= 0)> +-------------------------------------------------------------------------------- + +(toplevel + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_set + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal) + (bare_id) + (bare_id) + (integer_literal) + (bare_id) + (integer_literal) + (integer_literal) + (bare_id) + (bare_id) + (integer_literal) + (integer_literal)))))) + +================================================================================ +affine.min and affine.max with non-trivial affine maps +================================================================================ +func.func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) { + %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1] + %1 = affine.min affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2] + %2 = affine.min affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2] + %3 = affine.min affine_map<()[] -> (77, 78, 79)> ()[] + return +} + +func.func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) { + %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1] + %1 = affine.max affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2] + %2 = affine.max affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2] + %3 = affine.max affine_map<()[] -> (77, 78, 79)> ()[] + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (integer_literal) + (bare_id) + (integer_literal) + (bare_id)))) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal)))) + (value_use) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal)))) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (integer_literal) + (integer_literal) + (integer_literal))))))) + (operation + (custom_operation + (func_dialect)))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (integer_literal) + (bare_id) + (integer_literal) + (bare_id)))) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal)))) + (value_use) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal)))) + (value_use) + (value_use)))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (integer_literal) + (integer_literal) + (integer_literal))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +affine.dma_start and affine.dma_wait with complicated affine expressions +================================================================================ +func.func @test4(%arg0 : index, %arg1 : index) { + %0 = memref.alloc() : memref<100x100xf32> + %1 = memref.alloc() : memref<100x100xf32, 2> + %2 = memref.alloc() : memref<1xi32> + %c64 = arith.constant 64 : index + affine.for %i0 = 0 to 10 { + affine.for %i1 = 0 to 10 { + affine.dma_start %0[(%i0 + symbol(%arg0)) floordiv 3, %i1], + %1[%i0, (%i1 + symbol(%arg1)) mod 9 + 7], + %2[%i0 + %i1 + 11], %c64 + : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32> + affine.dma_wait %2[%i0 + %i1 + 11], %c64 : memref<1xi32> + } + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (integer_literal)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (integer_literal) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (integer_literal) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (value_use) + (value_use) + (integer_literal) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (value_use) + (value_use) + (value_use) + (integer_literal) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (integer_literal))))) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (affine_dialect + (value_use) + (value_use) + (value_use) + (integer_literal) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))))))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +affine_map with complex affine expression +================================================================================ +func.func @semiaffine_mod(%arg0: index, %arg1: index) -> index { + %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * s0) mod s0)> (%arg0)[%arg1] + return %a : index +} + +func.func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index { + %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + ((2 * s0) mod (3 * s0))) floordiv s0)> (%arg0)[%arg1] + return %a : index +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id)))) + (value_use) + (value_use)))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (index_type))))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal) + (bare_id) + (integer_literal) + (bare_id) + (bare_id)))) + (value_use) + (value_use)))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (index_type)))))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/attribute.txt b/mlir/utils/tree-sitter-mlir/test/corpus/attribute.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/attribute.txt @@ -0,0 +1,1301 @@ +================================================================================ +func.func with result attribute +================================================================================ +func.func @example_fn_result() -> (f64 {dialectName.attrName = 0 : i64}) +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (float_type))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (integer_literal) + (type + (builtin_type + (integer_type))))))))))))) + +================================================================================ +func.func with argument attribute +================================================================================ +func.func @example_fn_arg(%x: i32 {swift.self = unit}) +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (unit_literal)))))))))) + +================================================================================ +strided attribute in memref +================================================================================ +func.func @dot_view(%arg0: memref>, + %arg1: memref>, %arg2: memref) { + linalg.dot ins(%arg0, %arg1 : memref>, + memref>) + outs(%arg2: memref) + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (custom_operation + (linalg_dialect + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +dense attribute in linalg.depthwise_conv_1d_nwc_wcm +================================================================================ +func.func @depthwise_conv_1d_nwc_wcm(%input: tensor<1x12x8xf32>, %filter: tensor<3x8x8xf32>) + -> tensor<1x10x8x8xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<1x10x8x8xf32> + %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> + %0 = linalg.depthwise_conv_1d_nwc_wcm {dilations = dense<1> : tensor<1xi64>, + strides = dense<1> : tensor<1xi64>} + ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8x8xf32>) + outs(%fill : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> + return %0 : tensor<1x10x8x8xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (tensor_literal + (integer_literal)) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (attribute_entry + (bare_id) + (attribute_value + (tensor_literal + (integer_literal)) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))))) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +fastmath attribute in arith +================================================================================ +func.func @fastmath(%arg0: f32, %arg1: f32, %arg2: i32) { + %0 = arith.addf %arg0, %arg1 fastmath : f32 + %1 = arith.subf %arg0, %arg1 fastmath : f32 + %2 = arith.mulf %arg0, %arg1 fastmath : f32 + %3 = arith.divf %arg0, %arg1 fastmath : f32 + %4 = arith.remf %arg0, %arg1 fastmath : f32 + %5 = arith.negf %arg0 fastmath : f32 + %6 = arith.addf %arg0, %arg1 fastmath : f32 + %7 = arith.addf %arg0, %arg1 fastmath : f32 + %8 = arith.mulf %arg0, %arg1 fastmath : f32 + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (integer_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (fastmath_attr) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Generic function bracketed attributes +================================================================================ +"llvm.func"() ({ +}) {sym_name = "qux", function_type = !llvm.func, i64)>, + arg_attrs = [{llvm.noalias}, {}], xxx = {yyy = 42}} : () -> () +-------------------------------------------------------------------------------- + +(toplevel + (operation + (generic_operation + (string_literal) + (region) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (string_literal))) + (attribute_entry + (bare_id) + (attribute_value + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body + (pretty_dialect_item_body))))))) + (attribute_entry + (bare_id) + (attribute_value + (dictionary_attribute + (attribute_entry + (bare_id))) + (dictionary_attribute))) + (attribute_entry + (bare_id) + (attribute_value + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (integer_literal)))))))) + (function_type)))) + +================================================================================ +linalg.generic with attributes using attribute aliases +================================================================================ +#map0 = affine_map<(d0, d1) -> (d0, d1)> +#map1 = affine_map<(d0, d1) -> (d0)> +#map2 = affine_map<(d0) -> (d0)> + +func.func @add_broadcast_mul_fusion(%arg0: tensor, %arg1 : tensor, + %arg2 : tensor) -> tensor +{ + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %0 = tensor.dim %arg0, %c0 : tensor + %1 = tensor.empty(%0) : tensor + %2 = linalg.generic {indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]} + ins(%arg0, %arg1 : tensor, tensor) + outs(%1 : tensor) { + ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): + %3 = arith.addf %arg3, %arg4 : f32 + linalg.yield %3 : f32 + } -> tensor + %3 = tensor.dim %arg2, %c1 : tensor + %4 = tensor.empty(%0, %3) : tensor + %5 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} + ins(%2, %arg2 : tensor, tensor) + outs(%4 : tensor){ + ^bb0(%arg5: f32, %arg6: f32, %arg7: f32): + %6 = arith.mulf %arg5, %arg6 : f32 + linalg.yield %6 : f32 + } -> tensor + return %5 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (attribute_alias) + (attribute_alias) + (attribute_alias))) + (attribute_entry + (bare_id) + (attribute_value + (string_literal))))) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (attribute_alias) + (attribute_alias) + (attribute_alias))) + (attribute_entry + (bare_id) + (attribute_value + (string_literal) + (string_literal))))) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +linalg.generic with attribute alias +================================================================================ +#map0 = affine_map<() -> ()> + +#attrs = { + indexing_maps = [#map0, #map0, #map0], + iterator_types = [] +} + +func.func @main() -> (tensor) attributes {} { + %c0 = arith.constant 0 : i32 + %0 = tensor.from_elements %c0 : tensor + %c10 = arith.constant 10 : i32 + %1 = tensor.from_elements %c10 : tensor + cf.br ^bb1(%0 : tensor) + +^bb1(%2: tensor): // 2 preds: ^bb0, ^bb2 + %3 = tensor.empty() : tensor + %4 = linalg.generic #attrs + ins(%2, %1 : tensor, tensor) + outs(%3 : tensor) { + ^bb0(%arg0: i32, %arg1: i32, %arg2: i1): + %8 = arith.cmpi slt, %arg0, %arg1 : i32 + linalg.yield %8 : i1 + } -> tensor + %5 = tensor.extract %4[] : tensor + cf.cond_br %5, ^bb2(%2 : tensor), ^bb3(%2 : tensor) + +^bb2(%6: tensor): // pred: ^bb1 + %7 = tensor.empty() : tensor + %8 = linalg.generic #attrs + ins(%6, %6 : tensor, tensor) + outs(%7 : tensor) { + ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): + %9 = arith.addi %arg0, %arg1 : i32 + linalg.yield %9 : i32 + } -> tensor + cf.br ^bb3(%8 : tensor) + +^bb3(%10: tensor): // pred: ^bb1 + return %10 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map)))) + (attribute_alias_def + (attribute_value + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (attribute_alias) + (attribute_alias) + (attribute_alias))) + (attribute_entry + (bare_id) + (attribute_value))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (attribute + (dictionary_attribute)) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (cf_dialect + (successor + (caret_id) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (attribute_alias)) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (cf_dialect + (value_use) + (successor + (caret_id) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))) + (successor + (caret_id) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (attribute_alias)) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (cf_dialect + (successor + (caret_id) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (comment) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))))))))) + +================================================================================ +linalg.generic with inline attributes +================================================================================ +func.func @copy_view(%arg0: memref>, + %arg1: memref>) { + linalg.generic { + iterator_types = ["parallel"], + indexing_maps = [ affine_map<(i) -> (i)>, affine_map<(i) -> (i)>] } + ins(%arg0: memref>) + outs(%arg1: memref>) { + ^bb0(%a: f32, %b: f32): + linalg.yield %a : f32 + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))) + (region + (entry_block + (operation + (custom_operation + (linalg_dialect + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (string_literal))) + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id))) + (builtin_attribute + (affine_map + (bare_id) + (bare_id))))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +linalg.broadcast with inline attributes, without dictionary tokens +================================================================================ +func.func @broadcast(%input: tensor<8x32xf32>, + %init: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> { + %bcast = linalg.broadcast + ins(%input:tensor<8x32xf32>) + outs(%init:tensor<8x16x32xf32>) + dimensions = [1] + func.return %bcast : tensor<8x16x32xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (bare_attribute_entry + (bare_id) + (attribute_value + (integer_literal)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/controlflow.txt b/mlir/utils/tree-sitter-mlir/test/corpus/controlflow.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/controlflow.txt @@ -0,0 +1,2100 @@ +================================================================================ +Branching written with cf.br, cf.cond_br +================================================================================ +func.func @simple(i64, i1) -> i64 { +^bb0(%a: i64, %cond: i1): // Code dominated by ^bb0 may refer to %a + cf.cond_br %cond, ^bb1, ^bb2 + +^bb1: + cf.br ^bb3(%a: i64) // Branch passes %a as the argument + +^bb2: + %b = arith.addi %a, %a : i64 + cf.br ^bb3(%b: i64) // Branch passes %b as the argument + +// ^bb3 receives an argument, named %c, from predecessors +// and passes it on to bb4 along with %a. %a is referenced +// directly from its defining operation and is not passed through +// an argument of ^bb3. +^bb3(%c: i64): + cf.br ^bb4(%c, %a : i64, i64) + +^bb4(%d : i64, %e : i64): + %0 = arith.addi %d, %e : i64 + return %0 : i64 // Return is also a terminator. +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))))) + (comment) + (operation + (custom_operation + (cf_dialect + (value_use) + (successor + (caret_id)) + (successor + (caret_id)))))) + (block + (block_label + (caret_id)) + (operation + (custom_operation + (cf_dialect + (successor + (caret_id) + (value_use) + (type + (builtin_type + (integer_type)))))))) + (comment) + (block + (block_label + (caret_id)) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (cf_dialect + (successor + (caret_id) + (value_use) + (type + (builtin_type + (integer_type)))))))) + (comment) + (comment) + (comment) + (comment) + (comment) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (custom_operation + (cf_dialect + (successor + (caret_id) + (value_use) + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type)))))))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type))))))) + (comment)))))) + +================================================================================ +Switch statement written with cf.switch +================================================================================ +func.func @switch(%flag : i32, %caseOperand : i32) { + cf.switch %flag : i32, [ + default: ^bb1(%caseOperand : i32), + 42: ^bb2(%caseOperand : i32), + 43: ^bb3(%caseOperand : i32) + ] + + ^bb1(%bb1arg : i32): + return + ^bb2(%bb2arg : i32): + return + ^bb3(%bb3arg : i32): + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type)))) + (region + (entry_block + (operation + (custom_operation + (cf_dialect + (value_use) + (type + (builtin_type + (integer_type))) + (cf_case_label) + (successor + (caret_id) + (value_use) + (type + (builtin_type + (integer_type)))) + (cf_case_label + (integer_literal)) + (successor + (caret_id) + (value_use) + (type + (builtin_type + (integer_type)))) + (cf_case_label + (integer_literal)) + (successor + (caret_id) + (value_use) + (type + (builtin_type + (integer_type)))))))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (custom_operation + (func_dialect)))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (custom_operation + (func_dialect)))) + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Simple for-loop written with scf.for +================================================================================ +func.func @for_loop_with_increasing_arg() -> i1 { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + %c16 = arith.constant 16 : index + %0 = scf.for %arg0 = %c0 to %c4 step %c1 iter_args(%arg1 = %c0) -> index { + %10 = arith.addi %arg0, %arg1 : index + scf.yield %10 : index + } + %1 = arith.cmpi ule, %0, %c16 : index + func.return %1 : i1 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (index_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +Nested for-loop written using scf.for +================================================================================ +func.func @std_for(%arg0 : index, %arg1 : index, %arg2 : index) { + scf.for %i0 = %arg0 to %arg1 step %arg2 { + scf.for %i1 = %arg0 to %arg1 step %arg2 { + %min_cmp = arith.cmpi slt, %i0, %i1 : index + %min = arith.select %min_cmp, %i0, %i1 : index + %max_cmp = arith.cmpi sge, %i0, %i1 : index + %max = arith.select %max_cmp, %i0, %i1 : index + scf.for %i2 = %min to %max step %i1 { + } + } + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (region + (entry_block + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (region)))))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Conditional written with scf.if and scf.yield +================================================================================ +func.func @std_if_yield(%arg0: i1, %arg1: f32) +{ + %x, %y = scf.if %arg0 -> (f32, f32) { + %0 = arith.addf %arg1, %arg1 : f32 + %1 = arith.subf %arg1, %arg1 : f32 + scf.yield %0, %1 : f32, f32 + } else { + %0 = arith.subf %arg1, %arg1 : f32 + %1 = arith.addf %arg1, %arg1 : f32 + scf.yield %0, %1 : f32, f32 + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (float_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type)))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Switch statement using scf.index_switch +================================================================================ +func.func @switch(%arg0: index) -> i32 { + %0 = scf.index_switch %arg0 -> i32 + case 2 { + %c10_i32 = arith.constant 10 : i32 + scf.yield %c10_i32 : i32 + } + case 5 { + %c20_i32 = arith.constant 20 : i32 + scf.yield %c20_i32 : i32 + } + default { + %c30_i32 = arith.constant 30 : i32 + scf.yield %c30_i32 : i32 + } + + scf.index_switch %arg0 + default { + scf.yield + } + + return %0 : i32 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (integer_type))) + (scf_case_label + (integer_literal)) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))) + (scf_case_label + (integer_literal)) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))) + (scf_case_label) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (integer_type))))))))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (scf_case_label) + (region + (entry_block + (operation + (custom_operation + (scf_dialect)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +Parallel for-loop and reduction using scf.parallel and scf.reduce +================================================================================ +func.func @single_iteration_reduce(%A: index, %B: index) -> (index, index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %c6 = arith.constant 6 : index + %0:2 = scf.parallel (%i0, %i1) = (%c1, %c3) to (%c2, %c6) step (%c1, %c3) init(%A, %B) -> (index, index) { + scf.reduce(%i0) : index { + ^bb0(%lhs: index, %rhs: index): + %1 = arith.addi %lhs, %rhs : index + scf.reduce.return %1 : index + } + scf.reduce(%i1) : index { + ^bb0(%lhs: index, %rhs: index): + %2 = arith.muli %lhs, %rhs : index + scf.reduce.return %2 : index + } + scf.yield + } + return %0#0, %0#1 : index, index +} + +func.func @single_iteration_some(%A: memref) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %c6 = arith.constant 6 : index + %c7 = arith.constant 7 : index + %c10 = arith.constant 10 : index + scf.parallel (%i0, %i1, %i2) = (%c0, %c3, %c7) to (%c1, %c6, %c10) step (%c1, %c2, %c3) { + %c42 = arith.constant 42 : i32 + memref.store %c42, %A[%i0, %i1, %i2] : memref + scf.yield + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))) + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (index_type))) + (type + (builtin_type + (index_type))) + (region + (entry_block + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))))))))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))))))))))) + (operation + (custom_operation + (scf_dialect)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type))) + (type + (builtin_type + (index_type))))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (memref_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (scf_dialect)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +While loop written using scf.while and scf.condition +================================================================================ +func.func @while_cond_true() -> i1 { + %0 = scf.while () : () -> i1 { + %condition = "test.condition"() : () -> i1 + scf.condition(%condition) %condition : i1 + } do { + ^bb0(%arg0: i1): + "test.use"(%arg0) : (i1) -> () + scf.yield + } + return %0 : i1 +} + +func.func @while_unused_arg(%x : i32, %y : f64) -> i32 { + %0 = scf.while (%arg1 = %x, %arg2 = %y) : (i32, f64) -> (i32) { + %condition = "test.condition"(%arg1) : (i32) -> i1 + scf.condition(%condition) %arg1 : i32 + } do { + ^bb0(%arg1: i32): + %next = "test.use"(%arg1) : (i32) -> (i32) + scf.yield %next, %y : i32, f64 + } + return %0 : i32 +} + +func.func @infinite_while() { + %true = arith.constant true + scf.while : () -> () { + scf.condition(%true) + } do { + scf.yield + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (function_type + (type + (builtin_type + (integer_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (generic_operation + (string_literal) + (function_type + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (generic_operation + (string_literal) + (value_use) + (function_type + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type))))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (float_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (function_type + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (integer_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (generic_operation + (string_literal) + (value_use) + (function_type + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))))) + (operation + (op_result + (value_use)) + (generic_operation + (string_literal) + (value_use) + (function_type + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type))))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (bool_literal)))) + (operation + (custom_operation + (scf_dialect + (function_type) + (region + (entry_block + (operation + (custom_operation + (scf_dialect + (value_use)))))) + (region + (entry_block + (operation + (custom_operation + (scf_dialect)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +affine.for with attribute aliseses +================================================================================ +#map_simple0 = affine_map<()[] -> (10)> +#map_simple1 = affine_map<()[s0] -> (s0)> +#map_non_simple0 = affine_map<(d0)[] -> (d0)> +#map_non_simple1 = affine_map<(d0)[s0] -> (d0 + s0)> +#map_non_simple2 = affine_map<()[s0, s1] -> (s0 + s1)> +#map_non_simple3 = affine_map<()[s0] -> (s0 + 3)> +func.func @funcsimplemap(%arg0: index, %arg1: index) -> () { + affine.for %i0 = 0 to #map_simple0()[] { + affine.for %i1 = 0 to #map_simple1()[%arg1] { + affine.for %i2 = 0 to #map_non_simple0(%i0)[] { + affine.for %i3 = 0 to #map_non_simple1(%i0)[%arg1] { + affine.for %i4 = 0 to #map_non_simple2()[%arg1, %arg0] { + affine.for %i5 = 0 to #map_non_simple3()[%arg0] { + %c42_i32 = arith.constant 42 : i32 + } + } + } + } + } + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (integer_literal))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))) + (attribute_alias_def + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (integer_literal))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens)) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (value_use) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (value_use) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (value_use) + (value_use) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (value_use) + (value_use) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (integer_literal) + (attribute + (attribute_alias)) + (value_use) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))))))))))))))))))))))))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +affine.if and affine.yield +================================================================================ +func.func @affine_if() -> f32 { + %zero = arith.constant 0.0 : f32 + %0 = affine.if affine_set<() : ()> () -> f32 { + affine.yield %zero : f32 + } else { + affine.yield %zero : f32 + } + return %0 : f32 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (float_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (attribute + (builtin_attribute + (affine_set))) + (type + (builtin_type + (float_type))) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (float_type)))))))))))) + +================================================================================ +Parallel for loop written using affine.parallel +================================================================================ +func.func @parallel(%A : memref<100x100xf32>, %N : index) { + affine.parallel (%i0, %j0) = (0, 0) to (symbol(%N), 100) step (10, 10) { + %0:2 = affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) + reduce ("minf", "maxf") -> (f32, f32) { + %2 = affine.load %A[%i0 + %i0, %j0 + %j1] : memref<100x100xf32> + affine.yield %2, %2 : f32, f32 + } + } + return +} + +func.func @parallel_min_max(%a: index, %b: index, %c: index, %d: index) { + affine.parallel (%i, %j, %k) = (max(%a, %b), %b, max(%a, %c)) + to (%c, min(%c, %d), %b) { + affine.yield + } + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (integer_literal) + (value_use) + (integer_literal) + (string_literal) + (string_literal) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (affine_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (affine_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type)))))))))))))))) + (operation + (custom_operation + (func_dialect)))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (custom_operation + (affine_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (region + (entry_block + (operation + (custom_operation + (affine_dialect)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Execution of a region using scf.execute_region +================================================================================ +func.func @execute_region() -> i64 { + %res = scf.execute_region -> i64 { + %c1 = arith.constant 1 : i64 + scf.yield %c1 : i64 + } + + %res2:2 = scf.execute_region -> (i64, i64) { + %c1 = arith.constant 1 : i64 + scf.yield %c1, %c1 : i64, i64 + } + return %res : i64 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (type + (builtin_type + (integer_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (integer_type))))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type))))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +Parallel loop written using scf.forall and scf.forall.in_parallel +================================================================================ +func.func @normalized_forall(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c1 = arith.constant 1 : index + %num_threads = arith.constant 100 : index + %result = scf.forall (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> { + %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : + tensor<1xf32> into tensor<100xf32> + } + } + return +} + +func.func @explicit_loop_bounds_forall(%in: tensor<100xf32>, + %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %num_threads = arith.constant 100 : index + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> tensor<100xf32> { + %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : + tensor<1xf32> into tensor<100xf32> + } {mapping = [#gpu.thread]} + } + return +} + +func.func @normalized_forall_elide_terminator() -> () { + %num_threads = arith.constant 100 : index + scf.forall (%thread_idx) in (%num_threads) { + scf.forall.in_parallel { + } + } {mapping = [#gpu.thread]} + return + +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (scf_dialect + (region + (entry_block + (operation + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))))))) + (operation + (custom_operation + (func_dialect)))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (scf_dialect + (region + (entry_block + (operation + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (dialect_attribute + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body))))))))))))))) + (operation + (custom_operation + (func_dialect)))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens)) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (value_use) + (region + (entry_block + (operation + (custom_operation + (scf_dialect + (region)))))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (dialect_attribute + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body)))))))))) + (operation + (custom_operation + (func_dialect))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/func.txt b/mlir/utils/tree-sitter-mlir/test/corpus/func.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/func.txt @@ -0,0 +1,261 @@ +================================================================================ +Function prototype +================================================================================ +func.func @abort() +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list))))) + +================================================================================ +Simple function using func, arith dialects +================================================================================ +func.func @test_addi(%arg0 : i64, %arg1 : i64) -> i64 { + %0 = arith.addi %arg0, %arg1 : i64 + return %0 : i64 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +Function with multiple return values +================================================================================ +func.func @count(%x: i64) -> (i64, i64) { + return %x, %x: i64, i64 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +Variadic function +================================================================================ +llvm.func @variadic(...) +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (llvm_dialect + (symbol_ref_id) + (func_arg_list + (variadic)))))) + +================================================================================ +Variadic function with other arguments +================================================================================ +llvm.func @variadic_args(i32, i32, ...) +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (llvm_dialect + (symbol_ref_id) + (func_arg_list + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (integer_type))) + (variadic)))))) + +================================================================================ +Generic operation in a module, with attributes +================================================================================ +module { + "llvm.func"() ({ + }) {sym_name = "foo", function_type = !llvm.func} : () -> () +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (builtin_dialect + (region + (entry_block + (operation + (generic_operation + (string_literal) + (region) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (string_literal))) + (attribute_entry + (bare_id) + (attribute_value + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body)))))))) + (function_type))))))))) + +================================================================================ +Generic operation with successor and region +================================================================================ +func.func @terminator_with_regions() { + "region"()[^bb2] ({}) : () -> () +^bb2: + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (region + (entry_block + (operation + (generic_operation + (string_literal) + (successor + (caret_id)) + (region) + (function_type)))) + (block + (block_label + (caret_id)) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Function with private specifier, and func.call +================================================================================ +func.func private @source() -> tensor +func.func @call_source() -> tensor { + %0 = call @source() : () -> tensor + return %0 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (func_dialect + (symbol_ref_id) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/op.txt b/mlir/utils/tree-sitter-mlir/test/corpus/op.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/op.txt @@ -0,0 +1,3398 @@ +================================================================================ +linalg.map without ins +================================================================================ +func.func @map_no_inputs(%init: tensor<64xf32>) -> tensor<64xf32> { + %add = linalg.map + outs(%init:tensor<64xf32>) + () { + %0 = arith.constant 0.0: f32 + linalg.yield %0: f32 + } + func.return %add : tensor<64xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (block_arg_list) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +linalg.map with ins and outs +================================================================================ +func.func @map_binary_memref(%lhs: memref<64xf32>, %rhs: memref<64xf32>, + %init: memref<64xf32>) { + linalg.map + ins(%lhs, %rhs: memref<64xf32>, memref<64xf32>) + outs(%init:memref<64xf32>) + (%lhs_elem: f32, %rhs_elem: f32) { + %0 = arith.addf %lhs_elem, %rhs_elem: f32 + linalg.yield %0: f32 + } + func.return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (custom_operation + (linalg_dialect + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (block_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (float_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type))))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +linalg.generic with attribute dictionary +================================================================================ +func.func @gather_like( + %arg0 : tensor {bufferization.writable = false}, + %arg1 : tensor {bufferization.writable = false}, + %arg2 : tensor {bufferization.writable = true}) + -> tensor +{ + %0 = linalg.generic { + indexing_maps = [affine_map<(d0, d1) -> (d0)>, + affine_map<(d0, d1) -> (d0, d1)>], + iterator_types = ["parallel", "parallel"]} + ins(%arg1 : tensor) outs(%arg2 : tensor) { + ^bb0(%arg3: i32, %arg4 : f32): + %iv1 = linalg.index 1 : index + %1 = arith.index_cast %arg3: i32 to index + %2 = tensor.extract %arg0[%1, %iv1] : tensor + linalg.yield %2 : f32 + } -> tensor + return %0 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))) + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))) + (attribute_entry + (bare_id) + (attribute_value + (string_literal) + (string_literal))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (float_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (linalg_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +tensor.collapse_shape and tensor.expand_shape without reassocation +================================================================================ +func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor) + -> (tensor, tensor<1x1xf32>) { + %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor + %1 = tensor.expand_shape %0 [] : tensor into tensor<1x1xf32> + return %0, %1 : tensor, tensor<1x1xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (nested_idx_list) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (nested_idx_list) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +tensor.collapse_shape with reassocation +================================================================================ +func.func @legal_collapsing_reshape_dynamic_tensor + (%arg0: tensor) -> tensor +{ + %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] : + tensor into tensor + return %0 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal)) + (nested_idx_list + (integer_literal)) + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +tensor.gather and tensor.scatter +================================================================================ +func.func @gather_scatter(%dest : tensor<4x5x6xf32>, %indices: tensor<1x3x2xindex>, + %indices_i32: tensor<1x3x2xi32>) { + %gathered = tensor.gather %dest[%indices_i32] gather_dims([1, 2]) unique: + (tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<1x3x4x1x1xf32> + %rank_reduced_gathered = tensor.gather %dest[%indices] gather_dims([1, 2]) unique: + (tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<1x3x4xf32> + + %scattered = tensor.scatter %gathered into %dest[%indices] + scatter_dims([1, 2]) unique: + (tensor<1x3x4x1x1xf32>, tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<4x5x6xf32> + %rank_reduced_scattered = tensor.scatter %rank_reduced_gathered into %dest[%indices_i32] + scatter_dims([1, 2]) unique: + (tensor<1x3x4xf32>, tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<4x5x6xf32> + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (index_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (gather_dims_attr + (integer_literal) + (integer_literal)) + (unique_attr) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (gather_dims_attr + (integer_literal) + (integer_literal)) + (unique_attr) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (index_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (scatter_dims_attr + (integer_literal) + (integer_literal)) + (unique_attr) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (index_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (scatter_dims_attr + (integer_literal) + (integer_literal)) + (unique_attr) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.extract_slice +================================================================================ +func.func @slice(%t: tensor<8x16x4xf32>, %idx : index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %1 = tensor.extract_slice %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] + : tensor<8x16x4xf32> to tensor + %2 = tensor.extract_slice %t[0, 2, 0][4, 4, 4][1, 1, 1] + : tensor<8x16x4xf32> to tensor<4x4x4xf32> + %3 = tensor.extract_slice %t[0, 2, 0][4, 1, 4][1, 1, 1] + : tensor<8x16x4xf32> to tensor<4x4xf32> + + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.insert_slice +================================================================================ +func.func @insert_slice( + %t: tensor<8x16x4xf32>, + %td: tensor<8x?x4xf32>, + %t2: tensor<16x32x8xf32>, + %t3: tensor<4x4xf32>, + %idx : index, + %sz : index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %1 = tensor.insert_slice %t into %t2[%c0, %c0, %c0][8, 16, 4][%c1, %c1, %c1] + : tensor<8x16x4xf32> into tensor<16x32x8xf32> + %2 = tensor.insert_slice %t into %t2[%c0, %idx, %c0][8, 16, 4][%c1, 1, %c1] + : tensor<8x16x4xf32> into tensor<16x32x8xf32> + %3 = tensor.insert_slice %t3 into %t[0, 2, 0][4, 1, 4][1, 1, 1] + : tensor<4x4xf32> into tensor<8x16x4xf32> + %4 = tensor.insert_slice %td into %t[0, %idx, 0][8, %sz, 4][1, 1, 1] + : tensor<8x?x4xf32> into tensor<8x16x4xf32> + + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (value_use) + (integer_literal) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (integer_literal) + (value_use) + (integer_literal) + (integer_literal) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.rank, tensor.generate, tensor.yield, tensor.extract +================================================================================ +func.func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index { + %size = tensor.rank %tensor : tensor<*xf32> + %0 = tensor.generate %size { + ^bb0(%arg0: index): + %1 = tensor.dim %tensor, %arg0 : tensor<*xf32> + tensor.yield %1 : index + } : tensor + %1 = tensor.extract %0[%idx] : tensor + return %1 : index +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (index_type))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (index_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (index_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (index_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (index_type)))))))))))) + +================================================================================ +tensor.insert +================================================================================ +func.func @insert(%arg0: f32, %arg1: index, %arg2: tensor) { + %0 = tensor.insert %arg0 into %arg2[%arg1, %arg1, %arg1] : tensor + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (float_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.from_elements and tensor.reshape +================================================================================ +func.func @tensor.reshape(%t1: tensor) -> tensor<2x2x5xf32> { + %two = arith.constant 2 : i64 + %five = arith.constant 5 : i64 + %shape = tensor.from_elements %two, %two, %five : tensor<3xi64> + %reshaped = tensor.reshape %t1(%shape) : (tensor, tensor<3xi64>) -> tensor<2x2x5xf32> + return %reshaped : tensor<2x2x5xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +tensor.pad +================================================================================ +func.func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, + %pad_value: f32) -> tensor<6x?x?x?xf32> { + %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] { + ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): + tensor.yield %pad_value : f32 + } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32> + return %0 : tensor<6x?x?x?xf32> +} + +func.func @pad_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index) + -> tensor<5x6xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %0 = tensor.pad %arg0 nofold low[%a, 0] high[0, %a] { + ^bb0(%arg1: index, %arg2: index): + tensor.yield %cst : f32 + } : tensor<5x6xf32> to tensor<5x6xf32> + return %0 : tensor<5x6xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (float_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (integer_literal) + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (value_use) + (integer_literal) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))))) + (operation + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))))))) + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (nofold_attr) + (value_use) + (integer_literal) + (integer_literal) + (value_use) + (region + (block + (block_label + (caret_id) + (block_arg_list + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))))) + (operation + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (float_type)))))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +tensor.splat +================================================================================ +func.func @test_splat_op(%s : f32) { + %v = tensor.splat %s : tensor<8xf32> + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (float_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.reshape +================================================================================ +func.func @tensor_reshape(%unranked: tensor<*xf32>, %shape: tensor<1xi32>) { + %dyn_vec = tensor.reshape %unranked(%shape) + : (tensor<*xf32>, tensor<1xi32>) -> tensor + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +tensor.pack and tensor.unpack +================================================================================ +func.func @pack_nc_to_ncnc(%source: tensor<128x256xf32>, %dest: tensor<4x16x32x16xf32>) -> tensor<128x256xf32> { + %0 = tensor.pack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<128x256xf32> -> tensor<4x16x32x16xf32> + %1 = tensor.empty() : tensor<128x256xf32> + %2 = tensor.unpack %0 inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %1 : tensor<4x16x32x16xf32> -> tensor<128x256xf32> + return %2 : tensor<128x256xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (inner_dims_pos_attr + (integer_literal) + (integer_literal)) + (inner_tiles_attr + (integer_literal) + (integer_literal)) + (value_use) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (inner_dims_pos_attr + (integer_literal) + (integer_literal)) + (inner_tiles_attr + (integer_literal) + (integer_literal)) + (value_use) + (function_type + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +memref.alloc, memref.cast, and memref.copy +================================================================================ +func.func @memref_copy() { + %0 = memref.alloc() : memref<2xf32> + %1 = memref.cast %0 : memref<2xf32> to memref<*xf32> + %2 = memref.alloc() : memref<2xf32> + %3 = memref.cast %0 : memref<2xf32> to memref<*xf32> + memref.copy %1, %3 : memref<*xf32> to memref<*xf32> + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (memref_dialect + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +memref.expand_shape and memref.collapse_shape +================================================================================ +func.func @expand_collapse_shape_dynamic(%arg0: memref, + %arg1: memref>, + %arg2: memref>, + %arg3: memref>) { + %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : + memref into memref + %r0 = memref.expand_shape %0 [[0, 1], [2]] : + memref into memref + %1 = memref.collapse_shape %arg1 [[0, 1], [2]] : + memref> into + memref> + %r1 = memref.expand_shape %1 [[0, 1], [2]] : + memref> into + memref> + %2 = memref.collapse_shape %arg2 [[0, 1], [2]] : + memref> into + memref> + %r2 = memref.expand_shape %2 [[0, 1], [2]] : + memref> into + memref> + %3 = memref.collapse_shape %arg3 [[0, 1]] : + memref> into + memref> + %r3 = memref.expand_shape %3 [[0, 1]] : + memref> into memref + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal))))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal)))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal)))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal)))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout + (integer_literal))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +memref.realloc +================================================================================ +func.func @memref_realloc_sd(%src : memref<2xf32>, %d : index) -> memref{ + %0 = memref.realloc %src(%d) : memref<2xf32> to memref + return %0 : memref +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (memref_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +memref.prefetch +================================================================================ +func.func @prefetch_locality_hint(%i : index) { + %0 = memref.alloc() : memref<10xf32> + memref.prefetch %0[%i], read, locality<3>, data : memref<10xf32> + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (memref_dialect + (value_use) + (value_use) + (isWrite_attr) + (localityHint_attr + (integer_literal)) + (isDataCache_attr) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +memref.rank +================================================================================ +func.func @rank(%t : memref<4x4x?xf32>) { + %0 = memref.rank %t : memref<4x4x?xf32> + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +memref.store +================================================================================ +func.func @truncation_spillover(%arg0 : memref) -> index { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c49 = arith.constant 49 : index + %0 = scf.for %arg1 = %c0 to %c2 step %c1 iter_args(%arg2 = %c0) -> index { + %1 = arith.divsi %arg2, %c49 : index + %2 = arith.index_cast %1 : index to i32 + memref.store %2, %arg0[%c0] : memref + %3 = arith.addi %arg2, %arg1 : index + scf.yield %3 : index + } + func.return %0 : index +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (scf_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (index_type))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (type + (builtin_type + (index_type))) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (memref_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (scf_dialect + (value_use) + (type + (builtin_type + (index_type))))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (index_type)))))))))))) + +================================================================================ +memref.reinterpret_cast +================================================================================ +func.func @memref_reinterpret_cast_static_to_dynamic_sizes(%in: memref) + -> memref<10x?xf32, strided<[?, 1], offset: ?>> { + %out = memref.reinterpret_cast %in to + offset: [1], sizes: [10, 10], strides: [1, 1] + : memref to memref<10x?xf32, strided<[?, 1], offset: ?>> + return %out : memref<10x?xf32, strided<[?, 1], offset: ?>> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout)))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)) + (attribute_value + (builtin_attribute + (strided_layout))))))))))))))) + +================================================================================ +cf.assert +================================================================================ +func.func @assert(%arg : i1) { + cf.assert %arg, "Some message in case this assertion fails." + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type)))) + (region + (entry_block + (operation + (custom_operation + (cf_dialect + (value_use) + (string_literal)))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +vector.splat +================================================================================ +func.func @signExtendConstantSplat() -> vector<4xi16> { + %c-2 = arith.constant -2 : i8 + %splat = vector.splat %c-2 : vector<4xi8> + %ext = arith.extsi %splat : vector<4xi8> to vector<4xi16> + return %ext : vector<4xi16> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))))))))) + +================================================================================ +arith.cmpi with string parameter +================================================================================ +func.func @notCmpEQ(%arg0: i8, %arg1: i8) -> i1 { + %true = arith.constant true + %cmp = arith.cmpi "eq", %arg0, %arg1 : i8 + %ncmp = arith.xori %cmp, %true : i1 + return %ncmp : i1 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type))) + (value_use) + (type + (builtin_type + (integer_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (integer_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (bool_literal)))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (string_literal) + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (integer_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (integer_type)))))))))))) + +================================================================================ +bufferization.alloc_tensor +================================================================================ +func.func @test_alloc_tensor_op(%t: tensor, %sz: index) + -> tensor +{ + %0 = bufferization.alloc_tensor(%sz) : tensor + %1 = bufferization.alloc_tensor() copy(%t) : tensor + %2 = bufferization.alloc_tensor() : tensor<5x6xf32> + %3 = bufferization.alloc_tensor(%sz, %sz) : tensor + %4 = bufferization.alloc_tensor() copy(%t) {escape = true} : tensor + %5 = bufferization.alloc_tensor() copy(%t) {escape = false} : tensor + %c100 = arith.constant 100 : index + %6 = bufferization.alloc_tensor() size_hint=%c100 : tensor<100x100xf64, #CSR> + %7 = bufferization.alloc_tensor(%sz) {memory_space = "foo"} : tensor + return %1 : tensor +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)) + (tensor_encoding + (attribute_value + (attribute_alias))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (bufferization_dialect + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (string_literal))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))))))))) + +================================================================================ +vector.transfer_read and vector.transfer_write +================================================================================ +func.func @vector_transfer_ops(%arg0: memref, + %arg1 : memref>, + %arg2 : memref>, + %arg3 : memref>, + %arg4 : memref) { + %c3 = arith.constant 3 : index + %cst = arith.constant 3.0 : f32 + %f0 = arith.constant 0.0 : f32 + %c0 = arith.constant 0 : i32 + %i0 = arith.constant 0 : index + %i1 = arith.constant 1 : i1 + + %vf0 = vector.splat %f0 : vector<4x3xf32> + %v0 = vector.splat %c0 : vector<4x3xi32> + %vi0 = vector.splat %i0 : vector<4x3xindex> + %m = arith.constant dense<[0, 0, 1, 0, 1]> : vector<5xi1> + %m2 = vector.splat %i1 : vector<4x5xi1> + %0 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d0)>} : memref, vector<128xf32> + %1 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : memref, vector<3x7xf32> + %2 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d0)>} : memref, vector<128xf32> + %3 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d1)>} : memref, vector<128xf32> + %4 = vector.transfer_read %arg1[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref>, vector<1x1x4x3xf32> + %5 = vector.transfer_read %arg1[%c3, %c3], %vf0 {in_bounds = [false, true]} : memref>, vector<1x1x4x3xf32> + %6 = vector.transfer_read %arg2[%c3, %c3], %v0 : memref>, vector<5x24xi8> + %7 = vector.transfer_read %arg3[%c3, %c3], %vi0 : memref>, vector<5x48xi8> + %8 = vector.transfer_read %arg0[%c3, %c3], %f0, %m : memref, vector<5xf32> + %9 = vector.transfer_read %arg4[%c3, %c3, %c3], %f0, %m2 {permutation_map = affine_map<(d0, d1, d2)->(d1, d0, 0)>} : memref, vector<5x4x8xf32> + + vector.transfer_write %0, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0)>} : vector<128xf32>, memref + vector.transfer_write %1, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : vector<3x7xf32>, memref + vector.transfer_write %4, %arg1[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : vector<1x1x4x3xf32>, memref> + vector.transfer_write %5, %arg1[%c3, %c3] {in_bounds = [false, false]} : vector<1x1x4x3xf32>, memref> + vector.transfer_write %6, %arg2[%c3, %c3] : vector<5x24xi8>, memref> + vector.transfer_write %7, %arg3[%c3, %c3] : vector<5x48xi8>, memref> + vector.transfer_write %8, %arg0[%c3, %c3], %m : vector<5xf32>, memref + + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (float_type)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (integer_type)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (index_type)))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (index_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (float_type)))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal) + (bool_literal))))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (float_type)))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (integer_type)))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (index_type)))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (bare_id) + (integer_literal))))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (builtin_attribute + (affine_map + (bare_id) + (bare_id) + (bare_id) + (bare_id))))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (float_type))))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (bool_literal) + (bool_literal))))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (float_type))))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (integer_type))))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (memref_type + (dim_list + (vector_type + (vector_dim_list) + (index_type))))))))) + (operation + (custom_operation + (vector_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/corpus/type.txt b/mlir/utils/tree-sitter-mlir/test/corpus/type.txt new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/corpus/type.txt @@ -0,0 +1,1398 @@ +================================================================================ +Tensor type of known dimensions +================================================================================ +func.func @test_addi_tensor(%arg0 : tensor<8x8xi64>, %arg1 : tensor<8x8xi64>) -> tensor<8x8xi64> { + %0 = arith.addi %arg0, %arg1 : tensor<8x8xi64> + return %0 : tensor<8x8xi64> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))))))))) + +================================================================================ +vector type with one dimension +================================================================================ +func.func @test_addi_vector(%arg0 : vector<8xi64>, %arg1 : vector<8xi64>) -> vector<8xi64> { + %0 = arith.addi %arg0, %arg1 : vector<8xi64> + return %0 : vector<8xi64> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))))))))) + +================================================================================ +vector type without dimensions +================================================================================ +func.func @test_cmpi_vector_0d(%arg0 : vector, %arg1 : vector) -> vector { + %0 = arith.cmpi ult, %arg0, %arg1 : vector + return %0 : vector +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (vector_type + (integer_type)))) + (value_use) + (type + (builtin_type + (vector_type + (integer_type))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (integer_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (vector_type + (integer_type))))))))))))) + +================================================================================ +vector type with multiple dimensions +================================================================================ +func.func @cmpi_equal_vector_operands(%arg0: vector<1x8xi64>) + -> (vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, + vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, + vector<1x8xi1>, vector<1x8xi1>) { + %0 = arith.cmpi eq, %arg0, %arg0 : vector<1x8xi64> + %1 = arith.cmpi sle, %arg0, %arg0 : vector<1x8xi64> + %2 = arith.cmpi sge, %arg0, %arg0 : vector<1x8xi64> + %3 = arith.cmpi ule, %arg0, %arg0 : vector<1x8xi64> + %4 = arith.cmpi uge, %arg0, %arg0 : vector<1x8xi64> + %5 = arith.cmpi ne, %arg0, %arg0 : vector<1x8xi64> + %6 = arith.cmpi slt, %arg0, %arg0 : vector<1x8xi64> + %7 = arith.cmpi sgt, %arg0, %arg0 : vector<1x8xi64> + %8 = arith.cmpi ult, %arg0, %arg0 : vector<1x8xi64> + %9 = arith.cmpi ugt, %arg0, %arg0 : vector<1x8xi64> + return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9 + : vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, + vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, + vector<1x8xi1>, vector<1x8xi1> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))))))))) + +================================================================================ +Scalable vector type +================================================================================ +func.func @test_addi_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 : vector<[8]xi64>) + -> vector<[8]xi64> { + %0 = arith.addi %arg0, %arg1 : vector<[8]xi64> + return %0 : vector<[8]xi64> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (integer_type))))))))))))) + +================================================================================ +!llvm.ptr dialect type +================================================================================ +llvm.func @byvalattr(%arg0: !llvm.ptr {llvm.byval}) { + llvm.return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (llvm_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body)))) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id))))) + (region + (entry_block + (operation + (custom_operation + (llvm_dialect))))))))) + +================================================================================ +Linalg matmul with memref of unknown dimensions +================================================================================ +func.func @matmul(%arg0: memref, %M: index, %N: index, %K: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %A = memref.view %arg0[%c0][%M, %K] : memref to memref + %B = memref.view %arg0[%c0][%K, %N] : memref to memref + %C = memref.view %arg0[%c0][%M, %N] : memref to memref + linalg.matmul ins(%A, %B: memref, memref) + outs(%C: memref) + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type))))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type))) + (value_use) + (type + (builtin_type + (index_type)))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (memref_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (integer_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (linalg_dialect + (value_use) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (type + (builtin_type + (memref_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (memref_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Various tensors of dimensions `?` and `*` using tensor.cast +================================================================================ +func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor) { + %0 = tensor.cast %arg0 : tensor<*xf32> to tensor + %1 = tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32> + %2 = tensor.cast %arg2 : tensor to tensor<4x?xf32> + %3 = tensor.cast %2 : tensor<4x?xf32> to tensor + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type))))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +Dense and sparse matrix literals +================================================================================ +func.func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32, complex) { + %const_0 = arith.constant 0 : index + %const_1 = arith.constant 1 : index + %const_3 = arith.constant 3 : index + + %0 = arith.constant dense<4.0> : tensor<4xf32> + %ext_1 = tensor.extract %0[%arg0] : tensor<4xf32> + + // Fold an extract into a sparse with a sparse index. + %1 = arith.constant sparse<[[0, 0, 0], [1, 1, 1]], [-5.0, -2.0]> : tensor<4x4x4xf16> + %ext_2 = tensor.extract %1[%const_1, %const_1, %const_1] : tensor<4x4x4xf16> + + // Fold an extract into a sparse with a non sparse index. + %2 = arith.constant sparse<[[1, 1, 1]], [-2.0]> : tensor<2x2x2xf16> + %ext_3 = tensor.extract %2[%const_0, %const_0, %const_0] : tensor<2x2x2xf16> + + // Fold an extract into a dense tensor. + %3 = arith.constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]> : tensor<2x1x4xi32> + %ext_4 = tensor.extract %3[%const_1, %const_0, %const_3] : tensor<2x1x4xi32> + + // Fold an extract into a complex constant. + %4 = arith.constant dense<(1.2, 2.3)> : tensor> + %ext_5 = tensor.extract %4[] : tensor> + + return %ext_1, %ext_2, %ext_3, %ext_4, %ext_5 : f32, f16, f16, i32, complex +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (index_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (complex_type + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (index_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (float_literal)) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal)) + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal))) + (nested_idx_list + (float_literal) + (float_literal))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal))) + (nested_idx_list + (float_literal))) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (float_type)))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (nested_idx_list + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal))) + (nested_idx_list + (nested_idx_list + (integer_literal) + (integer_literal) + (integer_literal) + (integer_literal))))) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type)))))))) + (comment) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (tensor_literal + (complex_literal + (float_literal) + (float_literal))) + (type + (builtin_type + (tensor_type + (dim_list + (complex_type + (float_type))))))))) + (operation + (op_result + (value_use)) + (custom_operation + (tensor_dialect + (value_use) + (type + (builtin_type + (tensor_type + (dim_list + (complex_type + (float_type))))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (complex_type + (float_type))))))))))))) + +================================================================================ +Empty dense literal +================================================================================ +func.func @empty_dense_literal() -> () { + "foo2"(){bar = dense<> : tensor<0xi32>} : () -> () + return +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens)) + (region + (entry_block + (operation + (generic_operation + (string_literal) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (tensor_literal) + (type + (builtin_type + (tensor_type + (dim_list + (integer_type))))))))) + (function_type))) + (operation + (custom_operation + (func_dialect))))))))) + +================================================================================ +!llvm.struct dialect type +================================================================================ +llvm.func @return_s_long() -> + !llvm.struct<"long", (i32, struct<(i32, i1)>, f32, ptr>)> +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (llvm_dialect + (symbol_ref_id) + (func_arg_list) + (func_return + (type_list_attr_parens + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body + (pretty_dialect_item_body) + (pretty_dialect_item_body + (pretty_dialect_item_body)))))))))))) + +================================================================================ +Type alias definition +================================================================================ +!baz = i64 +!qux = !llvm.struct<(!baz)> +!rec = !llvm.struct<"a", (ptr>)> +llvm.func @aliases() { + "some.op"() : () -> !llvm.struct<(i32, f32, !qux)> + "some.op"() : () -> !rec + llvm.return +} +-------------------------------------------------------------------------------- + +(toplevel + (type_alias_def + (type + (builtin_type + (integer_type)))) + (type_alias_def + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body))))) + (type_alias_def + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body + (pretty_dialect_item_body + (pretty_dialect_item_body))))))) + (operation + (custom_operation + (llvm_dialect + (symbol_ref_id) + (func_arg_list) + (region + (entry_block + (operation + (generic_operation + (string_literal) + (function_type + (type + (dialect_type + (pretty_dialect_item + (dialect_namespace) + (dialect_ident) + (pretty_dialect_item_body))))))) + (operation + (generic_operation + (string_literal) + (function_type + (type + (type_alias))))) + (operation + (custom_operation + (llvm_dialect))))))))) + +================================================================================ +negative constant value_use +================================================================================ +func.func @indexCastFold(%arg0: i8) -> index { + %c-2 = arith.constant -2 : i8 + %idx = arith.index_cast %c-2 : i8 to index + return %idx : index +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (integer_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (index_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (integer_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (type + (builtin_type + (integer_type))) + (type + (builtin_type + (index_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (index_type)))))))))))) + +================================================================================ +inf value_use +================================================================================ +func.func @test_maxf(%arg0 : f32) -> (f32, f32, f32) { + %c0 = arith.constant 0.0 : f32 + %-inf = arith.constant 0xFF800000 : f32 + %0 = arith.maxf %c0, %arg0 : f32 + %1 = arith.maxf %arg0, %arg0 : f32 + %2 = arith.maxf %-inf, %arg0 : f32 + return %0, %1, %2 : f32, f32, f32 +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (float_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (float_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (integer_literal) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (op_result + (value_use)) + (custom_operation + (arith_dialect + (value_use) + (value_use) + (type + (builtin_type + (float_type)))))) + (operation + (custom_operation + (func_dialect + (value_use) + (value_use) + (value_use) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type))) + (type + (builtin_type + (float_type)))))))))))) + +================================================================================ +array literal +================================================================================ +func.func @foo(%a : f32) -> vector<1x2xf32> { + %0 = "test_create_broadcast"(%a) {broadcast_dims = array} : (f32) -> vector<1x2xf32> + return %0: vector<1x2xf32> +} +-------------------------------------------------------------------------------- + +(toplevel + (operation + (custom_operation + (func_dialect + (symbol_ref_id) + (func_arg_list + (value_use) + (type + (builtin_type + (float_type)))) + (func_return + (type_list_attr_parens + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type)))))) + (region + (entry_block + (operation + (op_result + (value_use)) + (generic_operation + (string_literal) + (value_use) + (attribute + (dictionary_attribute + (attribute_entry + (bare_id) + (attribute_value + (array_literal + (type + (builtin_type + (integer_type))) + (integer_literal) + (integer_literal)))))) + (function_type + (type + (builtin_type + (float_type))) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))) + (operation + (custom_operation + (func_dialect + (value_use) + (type + (builtin_type + (vector_type + (vector_dim_list) + (float_type))))))))))))) diff --git a/mlir/utils/tree-sitter-mlir/test/highlight/attribute.mlir b/mlir/utils/tree-sitter-mlir/test/highlight/attribute.mlir new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/highlight/attribute.mlir @@ -0,0 +1,117 @@ +func.func @depthwise_conv_1d_nwc_wcm(%input: tensor<1x12x8xf32>, %filter: tensor<3x8x8xf32>) +// <- function.builtin +// ^ function +// ^ variable.parameter +// ^ type.builtin +// ^ variable.parameter +// ^ type.builtin + -> tensor<1x10x8x8xf32> { +// ^ operator +// ^ type.builtin + %zero = arith.constant 0.000000e+00 : f32 +// ^ variable +// ^ function.builtin +// ^ number +// ^ type.builtin + %init = tensor.empty() : tensor<1x10x8x8xf32> +// ^ variable +// ^ function.builtin +// ^ type.builtin + %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> +// ^ variable +// ^ function.builtin +// ^ keyword +// ^ variable +// ^ type.builtin +// ^ keyword + %0 = linalg.depthwise_conv_1d_nwc_wcm {dilations = dense<1> : tensor<1xi64>, +// ^ variable +// ^ function.builtin +// ^ attribute +// ^ constant.builtin + strides = dense<1> : tensor<1xi64>} +// ^ constant.builtin + ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8x8xf32>) +// ^ variable.parameter +// ^ variable.parameter + outs(%fill : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> +// ^ variable + return %0 : tensor<1x10x8x8xf32> +// ^ function.builtin +// ^ variable +} + +func.func @fastmath(%arg0: f32, %arg1: f32) { +// <- function.builtin +// ^ function +// ^ variable.parameter +// ^ type.builtin +// ^ variable.parameter +// ^ type.builtin + %5 = arith.negf %arg0 fastmath : f32 +// ^ function.builtin +// ^ attribute + %6 = arith.addf %arg0, %arg1 fastmath : f32 +// ^ function.builtin +// ^ attribute + %8 = arith.mulf %arg0, %arg1 fastmath : f32 +// ^ function.builtin +// ^ attribute + return +// ^ function.builtin +} + +#map0 = affine_map<(d0, d1) -> (d0, d1)> +// <- attribute +// ^ attribute +#map1 = affine_map<(d0, d1) -> (d0)> +// <- attribute +// ^ attribute +#map2 = affine_map<(d0) -> (d0)> +// <- attribute +// ^ attribute + +func.func @add_broadcast_mul_fusion(%arg0: tensor, %arg1 : tensor, + %arg2 : tensor) -> tensor +{ + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %0 = tensor.dim %arg0, %c0 : tensor + %1 = tensor.empty(%0) : tensor + %2 = linalg.generic {indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]} +// ^ attribute +// ^ attribute +// ^ attribute + ins(%arg0, %arg1 : tensor, tensor) +// ^ keyword + outs(%1 : tensor) { +// ^ keyword + ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): + %3 = arith.addf %arg3, %arg4 : f32 + linalg.yield %3 : f32 + } -> tensor + %3 = tensor.dim %arg2, %c1 : tensor + %4 = tensor.empty(%0, %3) : tensor + %5 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} +// ^ function.builtin + ins(%2, %arg2 : tensor, tensor) + outs(%4 : tensor){ + ^bb0(%arg5: f32, %arg6: f32, %arg7: f32): + %6 = arith.mulf %arg5, %arg6 : f32 + linalg.yield %6 : f32 + } -> tensor + return %5 : tensor +} + +func.func @broadcast(%input: tensor<8x32xf32>, + %init: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> { + %bcast = linalg.broadcast +// ^ function.builtin + ins(%input:tensor<8x32xf32>) +// ^ keyword + outs(%init:tensor<8x16x32xf32>) +// ^ keyword + dimensions = [1] +// ^ attribute + func.return %bcast : tensor<8x16x32xf32> +} diff --git a/mlir/utils/tree-sitter-mlir/test/highlight/controlflow.mlir b/mlir/utils/tree-sitter-mlir/test/highlight/controlflow.mlir new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/highlight/controlflow.mlir @@ -0,0 +1,78 @@ +func.func @simple(i64, i1) -> i64 { +// <- function.builtin +// ^ function +// ^ punctuation.bracket +// ^ type.builtin +// ^ punctuation.delimeter +// ^ type.builtin +// ^ punctuation.bracket +// ^ operator +// ^ type.builtin +// ^ punctuation.bracket +^bb0(%a: i64, %cond: i1): +// <- tag +// ^ variable.parameter +// ^ type.builtin +// ^ variable.parameter +// ^ type.builtin + cf.cond_br %cond, ^bb1, ^bb2 +// ^ function.builtin +// ^ variable.parameter +// ^ tag +// ^ tag + +^bb1: +// <- tag + cf.br ^bb3(%a: i64) // Branch passes %a as the argument +// ^ function.builtin +// ^ tag +// ^ variable.parameter +// ^ type.builtin +// ^ comment + +^bb2: +// <- tag + %b = arith.addi %a, %a : i64 +// ^ variable +// ^ operator +// ^ function.builtin +// ^ variable.parameter +// ^ variable.parameter +// ^ type.builtin + cf.br ^bb3(%b: i64) // Branch passes %b as the argument +// ^ function.builtin +// ^ tag +// ^ variable +// ^ type.builtin +// ^ comment +^bb3(%c: i64): +// <- tag +// ^ variable.parameter +// ^ type.builtin + cf.br ^bb4(%c, %a : i64, i64) +// ^ function.builtin +// ^ tag +// ^ variable.parameter +// ^ variable.parameter +// ^ type.builtin +// ^ type.builtin +^bb4(%d : i64, %e : i64): +// <- tag +// ^ variable.parameter +// ^ type.builtin +// ^ variable.parameter +// ^ type.builtin + %0 = arith.addi %d, %e : i64 +// ^ variable +// ^ operator +// ^ function.builtin +// ^ variable.parameter +// ^ variable.parameter +// ^ type.builtin + return %0 : i64 // Return is also a terminator. +// ^ function.builtin +// ^ variable +// ^ type.builtin +// ^ comment +} +// <- punctuation.bracket diff --git a/mlir/utils/tree-sitter-mlir/test/highlight/downstream.mlir b/mlir/utils/tree-sitter-mlir/test/highlight/downstream.mlir new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/highlight/downstream.mlir @@ -0,0 +1,27 @@ +func.func @sort_memref(%input1: memref, %input2: memref, +// <- function.builtin +// ^ function +// ^ variable.parameter +// ^ type.builtin + %init1: memref, %init2: memref) { + thlo.sort + ins(%input1: memref, %input2: memref) +// ^ keyword +// ^ variable.parameter + outs(%init1: memref, %init2: memref) +// ^ keyword +// ^ variable.parameter + { dimension = 0 : i64, is_stable = true } +// ^ constant.builtin + (%e11: f32, %e12: f32, %e21: i32, %e22: i32) { + %gt = arith.cmpf ogt, %e11, %e12: f32 +// ^ function.builtin +// ^ keyword +// ^ variable +// ^ variable +// ^ type.builtin + thlo.yield %gt : i1 + } + func.return +// ^ function.builtin +} diff --git a/mlir/utils/tree-sitter-mlir/test/highlight/func.mlir b/mlir/utils/tree-sitter-mlir/test/highlight/func.mlir new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/highlight/func.mlir @@ -0,0 +1,27 @@ +func.func @test_addi(%arg0 : i64, %arg1 : i64) -> i64 { +// <- function.builtin +// ^ function +// ^ punctuation.bracket +// ^ variable.parameter +// ^ punctuation.delimeter +// ^ type.builtin +// ^ punctuation.delimeter +// ^ variable.parameter +// ^ type.builtin +// ^ punctuation.bracket +// ^ operator +// ^ type.builtin +// ^ punctuation.bracket + %0 = arith.addi %arg0, %arg1 : i64 +// ^ variable +// ^ operator +// ^ function.builtin +// ^ variable.parameter +// ^ variable.parameter +// ^ type.builtin + return %0 : i64 +// ^ function.builtin +// ^ variable +// ^ type.builtin +} +// <- punctuation.bracket diff --git a/mlir/utils/tree-sitter-mlir/test/highlight/type.mlir b/mlir/utils/tree-sitter-mlir/test/highlight/type.mlir new file mode 100644 --- /dev/null +++ b/mlir/utils/tree-sitter-mlir/test/highlight/type.mlir @@ -0,0 +1,19 @@ +!baz = i64 +// <- type +// ^ type.builtin +!qux = !llvm.struct<(!baz)> +// <- type +// ^ type +!rec = !llvm.struct<"a", (ptr>)> +// <- type +// ^ type +llvm.func @aliases() { +// <- function.builtin +// ^ function + "some.op"() : () -> !llvm.struct<(i32, f32, !qux)> +// ^ type + "some.op"() : () -> !rec +// ^ type + llvm.return +// ^ function.builtin +}