xref: /llvm-project/mlir/utils/tree-sitter-mlir/dialect/tensor.js (revision 10f8be19e7562624c0e63872e28cc963e03e56a5)
1'use strict';
2
3module.exports = {
4  tensor_dialect : $ => choice(
5                     // operation ::= `tensor.empty` `(`$dynamicSizes`)`
6                     // attr-dict `:` type($result)
7                     seq('tensor.empty',
8                         field('dynamicSizes', $._value_use_list_parens),
9                         field('attributes', optional($.attribute)),
10                         field('return', $._type_annotation)),
11
12                     // operation ::= `tensor.cast` $source attr-dict `:`
13                     // type($source) `to` type($dest)
14                     seq('tensor.cast', field('in', $.value_use),
15                         field('attributes', optional($.attribute)),
16                         field('return', $._type_annotation)),
17
18                     // operation ::= `tensor.dim` attr-dict $source `,` $index
19                     // `:` type($source)
20                     seq('tensor.dim',
21                         field('attributes', optional($.attribute)),
22                         field('tensor', $.value_use), ',',
23                         field('index', $.value_use),
24                         field('return', $._type_annotation)),
25
26                     // operation ::= `tensor.collapse_shape` $src
27                     // $reassociation attr-dict `:` type($src)
28                     //                `into` type($result)
29                     seq(choice('tensor.collapse_shape', 'tensor.expand_shape'),
30                         field('tensor', $.value_use),
31                         field('reassociation', $.nested_idx_list),
32                         field('attributes', optional($.attribute)),
33                         field('return', $._type_annotation)),
34
35                     // operation ::= `tensor.extract` $tensor `[` $indices `]`
36                     // attr-dict `:` type($tensor)
37                     seq('tensor.extract', field('tensor', $.value_use),
38                         field('indices', $._dense_idx_list),
39                         field('attributes', optional($.attribute)),
40                         field('return', $._type_annotation)),
41
42                     // operation ::= `tensor.insert` $scalar `into` $dest `[`
43                     // $indices `]` attr-dict
44                     //               `:` type($dest)
45                     seq('tensor.insert', field('scalar', $.value_use),
46                         token('into'), field('destination', $.value_use),
47                         field('indices', $._dense_idx_list),
48                         field('attributes', optional($.attribute)),
49                         field('return', $._type_annotation)),
50
51                     // operation ::= `tensor.extract_slice` $source ``
52                     //                custom<DynamicIndexList>($offsets,
53                     //                $static_offsets)
54                     //                custom<DynamicIndexList>($sizes,
55                     //                $static_sizes)
56                     //                custom<DynamicIndexList>($strides,
57                     //                $static_strides) attr-dict `:`
58                     //                type($source) `to` type($result)
59                     seq('tensor.extract_slice', field('tensor', $.value_use),
60                         field('offsets', $._dense_idx_list),
61                         field('sizes', $._dense_idx_list),
62                         field('strides', $._dense_idx_list),
63                         field('attributes', optional($.attribute)),
64                         field('return', $._type_annotation)),
65
66                     // operation ::= `tensor.insert_slice` $source `into` $dest
67                     // ``
68                     //                custom<DynamicIndexList>($offsets,
69                     //                $static_offsets)
70                     //                custom<DynamicIndexList>($sizes,
71                     //                $static_sizes)
72                     //                custom<DynamicIndexList>($strides,
73                     //                $static_strides) attr-dict `:`
74                     //                type($source) `into` type($dest)
75                     // operation ::= `tensor.parallel_insert_slice` $source
76                     // `into` $dest ``
77                     //                custom<DynamicIndexList>($offsets,
78                     //                $static_offsets)
79                     //                custom<DynamicIndexList>($sizes,
80                     //                $static_sizes)
81                     //                custom<DynamicIndexList>($strides,
82                     //                $static_strides) attr-dict `:`
83                     //                type($source) `into` type($dest)
84                     seq(choice('tensor.insert_slice',
85                                'tensor.parallel_insert_slice'),
86                         field('source', $.value_use), token('into'),
87                         field('destination', $.value_use),
88                         field('offsets', $._dense_idx_list),
89                         field('sizes', $._dense_idx_list),
90                         field('strides', $._dense_idx_list),
91                         field('attributes', optional($.attribute)),
92                         field('return', $._type_annotation)),
93
94                     // operation ::= `tensor.from_elements` $elements attr-dict
95                     // `:` type($result)
96                     seq('tensor.from_elements',
97                         field('elements', optional($._value_use_list)),
98                         field('attributes', optional($.attribute)),
99                         field('return', $._type_annotation)),
100
101                     // operation ::= `tensor.gather` $source `[` $indices `]`
102                     //               `gather_dims` `(` $gather_dims `)`
103                     //               (`unique` $unique^)?
104                     //               attr-dict
105                     //               `:` functional-type(operands, results)
106                     seq('tensor.gather', field('source', $.value_use),
107                         field('indices', $._dense_idx_list),
108                         field('gatherDims', $.gather_dims_attr),
109                         field('unique', optional($.unique_attr)),
110                         field('attributes', optional($.attribute)),
111                         field('return', $._function_type_annotation)),
112
113                     // operation ::= `tensor.scatter` $source `into` $dest `[`
114                     // $indices `]`
115                     //               `scatter_dims` `(` $scatter_dims `)`
116                     //               (`unique` $unique^)?
117                     //               attr-dict
118                     //               `:` functional-type(operands, results)
119                     seq('tensor.scatter', field('source', $.value_use),
120                         token('into'), field('destination', $.value_use),
121                         field('indices', $._dense_idx_list),
122                         field('scatterDims', $.scatter_dims_attr),
123                         field('unique', optional($.unique_attr)),
124                         field('attributes', optional($.attribute)),
125                         field('return', $._function_type_annotation)),
126
127                     // operation ::= `tensor.pad` $source
128                     //               (`nofold` $nofold^)?
129                     //               `low` `` custom<DynamicIndexList>($low,
130                     //               $static_low) `high` ``
131                     //               custom<DynamicIndexList>($high,
132                     //               $static_high) $region attr-dict `:`
133                     //               type($source) `to` type($result)
134                     seq('tensor.pad', field('source', $.value_use),
135                         field('nofold', optional($.nofold_attr)),
136                         field('low', seq(token('low'), $._dense_idx_list)),
137                         field('high', seq(token('high'), $._dense_idx_list)),
138                         field('body', $.region),
139                         field('attributes', optional($.attribute)),
140                         field('return', $._type_annotation)),
141
142                     // operation ::= `tensor.reshape` $source `(` $shape `)`
143                     // attr-dict
144                     //                `:` functional-type(operands, results)
145                     seq('tensor.reshape', field('tensor', $.value_use),
146                         field('shape', $._value_use_list_parens),
147                         field('attributes', optional($.attribute)),
148                         field('return', $._function_type_annotation)),
149
150                     // operation ::= `tensor.splat` $input attr-dict `:`
151                     // type($aggregate)
152                     seq('tensor.splat', field('input', $.value_use),
153                         field('attributes', optional($.attribute)),
154                         field('return', $._type_annotation)),
155
156                     // operation ::= `tensor.pack` $source
157                     //               (`padding_value` `(` $padding_value^ `:`
158                     //               type($padding_value) `)`)?
159                     //               (`outer_dims_perm` `=` $outer_dims_perm^)?
160                     //               `inner_dims_pos` `=` $inner_dims_pos
161                     //               `inner_tiles` `=`
162                     //               custom<DynamicIndexList>($inner_tiles,
163                     //               $static_inner_tiles) `into` $dest
164                     //               attr-dict `:` type($source) `->`
165                     //               type($dest)
166                     // operation ::= `tensor.unpack` $source
167                     //               (`outer_dims_perm` `=` $outer_dims_perm^)?
168                     //               `inner_dims_pos` `=` $inner_dims_pos
169                     //               `inner_tiles` `=`
170                     //               custom<DynamicIndexList>($inner_tiles,
171                     //               $static_inner_tiles) `into` $dest
172                     //               attr-dict `:` type($source) `->`
173                     //               type($dest)
174                     seq(choice('tensor.pack', 'tensor.unpack'),
175                         field('source', $.value_use),
176                         field('padding_value',
177                               optional(seq(token('padding_value'), '(',
178                                            $._value_use_and_type, ')'))),
179                         field('outer_dims_perm',
180                               optional($.outer_dims_perm_attr)),
181                         field('inner_dims_pos', $.inner_dims_pos_attr),
182                         field('inner_tiles', $.inner_tiles_attr),
183                         token('into'), field('destination', $.value_use),
184                         field('return', $._function_type_annotation)),
185
186                     // operation ::= `tensor.generate` $dynamicExtents $body
187                     // attr-dict `:` type($result)
188                     seq('tensor.generate',
189                         field('dynamicExtents', $._value_use_list),
190                         field('body', $.region),
191                         field('attributes', optional($.attribute)),
192                         field('return', $._type_annotation)),
193
194                     // operation ::= `tensor.rank` $tensor attr-dict `:`
195                     // type($tensor) operation ::= `tensor.yield` $value
196                     // attr-dict `:` type($value)
197                     seq(choice('tensor.rank', 'tensor.yield'),
198                         field('tensor', $.value_use),
199                         field('attributes', optional($.attribute)),
200                         field('return', $._type_annotation)))
201}
202