Skip to content

Commit 10f8be1

Browse files
artagnonartagnon
committed
tree-sitter-mlir: add a more complete grammar
Contribute a grammar, along with associated tests, from the upstream project maintained at https://github.com/artagnon/tree-sitter-mlir. The new grammar includes several fixes, and successfully parses 60-80% of MLIR tests in the Arith, Math, ControlFlow, SCF, Tensor, Affine, and Linalg dialects. Differential Revision: https://reviews.llvm.org/D144408
1 parent 91be60b commit 10f8be1

31 files changed

+11378
-265
lines changed

‎mlir/utils/tree-sitter-mlir/.gitignore

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
node_modules/
2+
src/
3+
bindings/
4+
binding.gyp
5+
Cargo.toml
6+
Cargo.lock
7+
target/
8+
build/

‎mlir/utils/tree-sitter-mlir/README.md

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,18 @@
1-
tree-sitter-mlir
2-
================
1+
# tree-sitter-mlir
32

4-
Basic [tree-sitter](https://github.com/tree-sitter/tree-sitter) grammar for
5-
MLIR following the [lang-ref](https://mlir.llvm.org/docs/LangRef/).
3+
[tree-sitter](https://github.com/tree-sitter/tree-sitter) grammar for MLIR
4+
following the [lang-ref](https://mlir.llvm.org/docs/LangRef/). The parser is
5+
incomplete, and the bench statistics on the test files in the MLIR tree are as
6+
follows:
67

7-
Note: the directory in [LLVM repo](https://github.com/llvm/llvm-project/)
8-
merely contains the grammar file(s) and not the NPM/generated code.
8+
```
9+
Math, 100% passed
10+
Builtin, 100% passed
11+
Func, 100% passed
12+
ControlFlow, 100% passed
13+
Tensor, 93.33% passed
14+
Arith, 83.33% passed
15+
SCF, 88% passed
16+
Affine, 73.08% passed
17+
Linalg, 51.11% passed
18+
```

‎mlir/utils/tree-sitter-mlir/dialect/affine.js

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
'use strict';
2+
3+
module.exports = {
4+
affine_dialect : $ => prec.right(choice(
5+
seq('affine.apply',
6+
field('operand',
7+
seq($.attribute, $._dim_and_symbol_use_list)),
8+
field('attributes', optional($.attribute))),
9+
10+
// operation ::= `affine.delinearize_index` $linear_index
11+
// `into` ` `
12+
// `(` $basis `)` attr-dict `:`
13+
// type($multi_index)
14+
seq('affine.delinearlize_index',
15+
field('operand', $.value_use), 'into',
16+
field('basis', $._value_use_list_parens),
17+
field('attributes', optional($.attribute)),
18+
field('return', $._type_annotation)),
19+
20+
// operation ::= `affine.dma_start` ssa-use `[`
21+
// multi-dim-affine-map-of-ssa-ids `]`,
22+
// `[` multi-dim-affine-map-of-ssa-ids `]`,
23+
// `[` multi-dim-affine-map-of-ssa-ids `]`,
24+
// ssa-use `:` memref-type
25+
seq(choice('affine.dma_start', 'affine.dma_wait'),
26+
field('operands',
27+
seq($.value_use, $._multi_dim_affine_expr_sq,
28+
repeat(seq(',', $.value_use,
29+
$._multi_dim_affine_expr_sq)))),
30+
',', field('numElements', $._value_use_list),
31+
field('return', $._type_annotation)),
32+
33+
// operation ::= `affine.for` ssa-id `=` lower-bound `to`
34+
// upper-bound
35+
// (`step` integer-literal)? `{` op* `}`
36+
seq('affine.for', field('iv', $.value_use), '=',
37+
field('lowerBound',
38+
seq(optional(token('max')), $._bound)),
39+
token('to'),
40+
field('upperBound',
41+
seq(optional(token('min')), $._bound)),
42+
field('step',
43+
optional(seq(token('step'), $.integer_literal))),
44+
field('iter_args',
45+
optional(seq(token('iter_args'),
46+
$._value_assignment_list))),
47+
field('return', optional($._function_return)),
48+
field('body', $.region),
49+
field('attributes', optional($.attribute))),
50+
51+
// operation ::= `affine.if` if-op-cond `{` op* `}`
52+
// (`else` `{` op* `}`)? if-op-cond ::= integer-set-attr
53+
// dim-and-symbol-use-list
54+
seq('affine.if',
55+
field('condition',
56+
seq($.attribute, $._dim_and_symbol_use_list)),
57+
field('return', optional($._function_return)),
58+
field('trueblk', $.region),
59+
field('falseblk',
60+
optional(seq(token('else'), $.region))),
61+
field('attributes', optional($.attribute))),
62+
63+
// operation ::= `affine.load` ssa-use `[`
64+
// multi-dim-affine-map-of-ssa-ids `]`
65+
// `:` memref-type
66+
seq(choice('affine.load', 'affine.vector_load'),
67+
field('operand', $.value_use),
68+
field('multiDimAffineMap',
69+
$._multi_dim_affine_expr_sq),
70+
field('return', $._type_annotation)),
71+
72+
// operation ::= `affine.min` affine-map-attribute
73+
// dim-and-symbol-use-list
74+
seq(choice('affine.min', 'affine.max'),
75+
field('operand',
76+
seq($.attribute, $._dim_and_symbol_use_list))),
77+
78+
seq('affine.parallel',
79+
field('iv', $._value_use_list_parens), '=',
80+
field('lowerBound', $._multi_dim_affine_expr_parens),
81+
token('to'),
82+
field('upperBound', $._multi_dim_affine_expr_parens),
83+
field('step',
84+
optional(seq(token('step'),
85+
$._multi_dim_affine_expr_parens))),
86+
field('reduce',
87+
optional(seq(token('reduce'), '(',
88+
$.string_literal,
89+
repeat(seq(',', $.string_literal)),
90+
')'))),
91+
field('return', optional($._function_return)),
92+
field('body', $.region)),
93+
94+
seq('affine.prefetch', field('source', $.value_use),
95+
field('indices',
96+
optional($._multi_dim_affine_expr_sq)),
97+
',', field('isWrite', $.isWrite_attr), ',',
98+
field('localityHint', $.localityHint_attr), ',',
99+
field('isDataCache', $.isDataCache_attr),
100+
field('attributes', optional($.attribute)),
101+
field('return', $._type_annotation)),
102+
103+
// operation ::= `affine.store` ssa-use, ssa-use `[`
104+
// multi-dim-affine-map-of-ssa-ids `]`
105+
// `:` memref-type
106+
seq(choice('affine.store', 'affine.vector_store'),
107+
field('source', $.value_use), ',',
108+
field('destination', $.value_use),
109+
field('multiDimAffineMap',
110+
$._multi_dim_affine_expr_sq),
111+
field('return', $._type_annotation)),
112+
113+
// operation ::= `affine.yield` attr-dict ($operands^ `:`
114+
// type($operands))?
115+
seq('affine.yield',
116+
field('attributes', optional($.attribute)),
117+
field('results', optional($._value_use_type_list)))))
118+
}

‎mlir/utils/tree-sitter-mlir/dialect/arith.js

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
'use strict';
2+
3+
module.exports = {
4+
arith_dialect : $ => choice(
5+
// operation ::= `arith.constant` attr-dict $value
6+
seq('arith.constant',
7+
field('attributes', optional($.attribute)),
8+
field('value', $._literal_and_type)),
9+
10+
// operation ::= `arith.addi` $lhs `,` $rhs attr-dict `:`
11+
// type($result) operation ::= `arith.subi` $lhs `,` $rhs
12+
// attr-dict `:` type($result) operation ::= `arith.divsi`
13+
// $lhs `,` $rhs attr-dict `:` type($result) operation ::=
14+
// `arith.divui` $lhs `,` $rhs attr-dict `:` type($result)
15+
// operation ::= `arith.ceildivsi` $lhs `,` $rhs attr-dict
16+
// `:` type($result) operation ::= `arith.ceildivui` $lhs
17+
// `,` $rhs attr-dict `:` type($result) operation ::=
18+
// `arith.floordivsi` $lhs `,` $rhs attr-dict `:`
19+
// type($result) operation ::= `arith.remsi` $lhs `,` $rhs
20+
// attr-dict `:` type($result) operation ::= `arith.remui`
21+
// $lhs `,` $rhs attr-dict `:` type($result) operation ::=
22+
// `arith.muli` $lhs `,` $rhs attr-dict `:` type($result)
23+
// operation ::= `arith.mulsi_extended` $lhs `,` $rhs
24+
// attr-dict `:` type($lhs) operation ::=
25+
// `arith.mului_extended` $lhs `,` $rhs attr-dict `:`
26+
// type($lhs) operation ::= `arith.andi` $lhs `,` $rhs
27+
// attr-dict `:` type($result) operation ::= `arith.ori`
28+
// $lhs `,` $rhs attr-dict `:` type($result) operation ::=
29+
// `arith.xori` $lhs `,` $rhs attr-dict `:` type($result)
30+
// operation ::= `arith.maxsi` $lhs `,` $rhs attr-dict `:`
31+
// type($result) operation ::= `arith.maxui` $lhs `,` $rhs
32+
// attr-dict `:` type($result) operation ::= `arith.minsi`
33+
// $lhs `,` $rhs attr-dict `:` type($result) operation ::=
34+
// `arith.minui` $lhs `,` $rhs attr-dict `:` type($result)
35+
// operation ::= `arith.shli` $lhs `,` $rhs attr-dict `:`
36+
// type($result) operation ::= `arith.shrsi` $lhs `,` $rhs
37+
// attr-dict `:` type($result) operation ::= `arith.shrui`
38+
// $lhs `,` $rhs attr-dict `:` type($result)
39+
seq(choice('arith.addi', 'arith.subi', 'arith.divsi',
40+
'arith.divui', 'arith.ceildivsi',
41+
'arith.ceildivui', 'arith.floordivsi',
42+
'arith.remsi', 'arith.remui', 'arith.muli',
43+
'arith.mulsi_extended', 'arith.mului_extended',
44+
'arith.andi', 'arith.ori', 'arith.xori',
45+
'arith.maxsi', 'arith.maxui', 'arith.minsi',
46+
'arith.minui', 'arith.shli', 'arith.shrsi',
47+
'arith.shrui'),
48+
field('lhs', $.value_use), ',',
49+
field('rhs', $.value_use),
50+
field('attributes', optional($.attribute)),
51+
field('return', $._type_annotation)),
52+
53+
// operation ::= `arith.addui_extended` $lhs `,` $rhs
54+
// attr-dict `:` type($sum)
55+
// `,` type($overflow)
56+
seq('arith.addui_extended', field('lhs', $.value_use), ',',
57+
field('rhs', $.value_use),
58+
field('attributes', optional($.attribute)),
59+
field('return', seq(':', $.type, ',', $.type))),
60+
61+
// operation ::= `arith.addf` $lhs `,` $rhs (`fastmath` ``
62+
// $fastmath^)?
63+
// attr-dict `:` type($result)
64+
// operation ::= `arith.divf` $lhs `,` $rhs (`fastmath` ``
65+
// $fastmath^)?
66+
// attr-dict `:` type($result)
67+
// operation ::= `arith.maxf` $lhs `,` $rhs (`fastmath` ``
68+
// $fastmath^)?
69+
// attr-dict `:` type($result)
70+
// operation ::= `arith.minf` $lhs `,` $rhs (`fastmath` ``
71+
// $fastmath^)?
72+
// attr-dict `:` type($result)
73+
// operation ::= `arith.mulf` $lhs `,` $rhs (`fastmath` ``
74+
// $fastmath^)?
75+
// attr-dict `:` type($result)
76+
// operation ::= `arith.remf` $lhs `,` $rhs (`fastmath` ``
77+
// $fastmath^)?
78+
// attr-dict `:` type($result)
79+
// operation ::= `arith.subf` $lhs `,` $rhs (`fastmath` ``
80+
// $fastmath^)?
81+
// attr-dict `:` type($result)
82+
seq(choice('arith.addf', 'arith.divf', 'arith.maxf',
83+
'arith.minf', 'arith.mulf', 'arith.remf',
84+
'arith.subf'),
85+
field('lhs', $.value_use), ',',
86+
field('rhs', $.value_use),
87+
field('fastmath', optional($.fastmath_attr)),
88+
field('attributes', optional($.attribute)),
89+
field('return', $._type_annotation)),
90+
91+
// operation ::= `arith.negf` $operand (`fastmath` ``
92+
// $fastmath^)?
93+
// attr-dict `:` type($result)
94+
seq(choice('arith.negf'), field('operand', $.value_use),
95+
field('fastmath', optional($.fastmath_attr)),
96+
field('attributes', optional($.attribute)),
97+
field('return', $._type_annotation)),
98+
99+
// operation ::= `arith.cmpi` $predicate `,` $lhs `,` $rhs
100+
// attr-dict `:` type($lhs) operation ::= `arith.cmpf`
101+
// $predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)
102+
seq(choice('arith.cmpi', 'arith.cmpf'),
103+
field('predicate',
104+
choice('eq', 'ne', 'oeq', 'olt', 'ole', 'ogt',
105+
'oge', 'slt', 'sle', 'sgt', 'sge', 'ult',
106+
'ule', 'ugt', 'uge', $.string_literal)),
107+
',', field('lhs', $.value_use), ',',
108+
field('rhs', $.value_use),
109+
field('attributes', optional($.attribute)),
110+
field('return', $._type_annotation)),
111+
112+
// operation ::= `arith.extf` $in attr-dict `:` type($in)
113+
// `to` type($out) operation ::= `arith.extsi` $in attr-dict
114+
// `:` type($in) `to` type($out) operation ::= `arith.extui`
115+
// $in attr-dict `:` type($in) `to` type($out) operation ::=
116+
// `arith.fptosi` $in attr-dict `:` type($in) `to`
117+
// type($out) operation ::= `arith.fptoui` $in attr-dict `:`
118+
// type($in) `to` type($out) operation ::=
119+
// `arith.index_cast` $in attr-dict `:` type($in) `to`
120+
// type($out) operation ::= `arith.index_castui` $in
121+
// attr-dict `:` type($in) `to` type($out) operation ::=
122+
// `arith.sitofp` $in attr-dict `:` type($in) `to`
123+
// type($out) operation ::= `arith.uitofp` $in attr-dict `:`
124+
// type($in) `to` type($out) operation ::= `arith.bitcast`
125+
// $in attr-dict `:` type($in) `to` type($out) operation ::=
126+
// `arith.truncf` $in attr-dict `:` type($in) `to`
127+
// type($out) operation ::= `arith.trunci` $in attr-dict `:`
128+
// type($in) `to` type($out)
129+
seq(choice('arith.extf', 'arith.extsi', 'arith.extui',
130+
'arith.fptosi', 'arith.fptoui',
131+
'arith.index_cast', 'arith.index_castui',
132+
'arith.sitofp', 'arith.uitofp', 'arith.bitcast',
133+
'arith.truncf', 'arith.trunci'),
134+
field('in', $.value_use),
135+
field('attributes', optional($.attribute)),
136+
field('return', $._type_annotation)),
137+
138+
seq('arith.select', field('cond', $.value_use), ',',
139+
field('trueblk', $.value_use), ',',
140+
field('falseblk', $.value_use),
141+
field('return', $._type_annotation)))
142+
}

‎mlir/utils/tree-sitter-mlir/dialect/bufferization.js

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
'use strict';
2+
3+
module.exports = {
4+
bufferization_dialect : $ => choice(
5+
seq('bufferization.alloc_tensor',
6+
field('in', $._value_use_list_parens),
7+
field('copy', optional(seq(token('copy'), '(',
8+
$.value_use, ')'))),
9+
field('size_hint',
10+
optional(seq(token('size_hint'), '=',
11+
$.value_use))),
12+
field('attributes', optional($.attribute)),
13+
field('return', $._type_annotation)),
14+
15+
// operation ::= `bufferization.to_memref` $tensor
16+
// attr-dict `:` type($memref)
17+
seq('bufferization.to_memref',
18+
field('tensor', $.value_use),
19+
field('attributes', optional($.attribute)),
20+
field('return', $._type_annotation)),
21+
22+
// operation ::= `bufferization.to_tensor` $memref
23+
// (`restrict` $restrict^)?
24+
// (`writable` $writable^)? attr-dict
25+
// `:` type($memref)
26+
seq('bufferization.to_tensor',
27+
field('memref', $.value_use),
28+
field('restrict', optional($.restrict_attr)),
29+
field('writable', optional($.writable_attr)),
30+
field('attributes', optional($.attribute)),
31+
field('return', $._type_annotation)))
32+
}

‎mlir/utils/tree-sitter-mlir/dialect/builtin.js

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
'use strict';
2+
3+
module.exports = {
4+
builtin_dialect : $ => prec.right(choice(
5+
// operation ::= `builtin.module` ($sym_name^)?
6+
// attr-dict-with-keyword $bodyRegion
7+
seq(choice('builtin.module', 'module'),
8+
field('name', optional($.bare_id)),
9+
field('attributes', optional($.attribute)),
10+
field('body', $.region)),
11+
12+
// operation ::= `builtin.unrealized_conversion_cast`
13+
// ($inputs^ `:` type($inputs))?
14+
// `to` type($outputs) attr-dict
15+
seq(choice('builtin.unrealized_conversion_cast',
16+
'unrealized_conversion_cast'),
17+
field('inputs', optional($._value_use_type_list)),
18+
token('to'), field('outputs', $._type_list_no_parens),
19+
field('attributes', optional($.attribute)))))
20+
}

0 commit comments

Comments
 (0)