xref: /llvm-project/mlir/utils/tree-sitter-mlir/dialect/vector.js (revision 10f8be19e7562624c0e63872e28cc963e03e56a5)
1'use strict';
2
3module.exports = {
4  vector_dialect : $ => prec.right(choice(
5                     // operation ::= `vector.bitcast` $source attr-dict `:`
6                     // type($source) `to` type($result) operation ::=
7                     // `vector.broadcast` $source attr-dict `:` type($source)
8                     // `to` type($vector) operation ::= `vector.shape_cast`
9                     // $source attr-dict `:` type($source) `to` type($result)
10                     // operation ::= `vector.type_cast` $memref attr-dict `:`
11                     // type($memref) `to` type($result)
12                     seq(choice('vector.bitcast', 'vector.broadcast',
13                                'vector.shape_cast', 'vector.type_cast'),
14                         field('in', $.value_use),
15                         field('attributes', optional($.attribute)),
16                         field('return', $._type_annotation)),
17
18                     // operation ::= `vector.constant_mask` $mask_dim_sizes
19                     // attr-dict `:` type(results)
20                     seq('vector.constant_mask',
21                         field('mask', $._dense_idx_list),
22                         field('attributes', optional($.attribute)),
23                         field('return', $._type_annotation)),
24
25                     // operation ::= `vector.create_mask` $operands attr-dict
26                     // `:` type(results)
27                     seq('vector.create_mask',
28                         field('operands', $._value_use_list),
29                         field('attributes', optional($.attribute)),
30                         field('return', $._type_annotation)),
31
32                     // operation ::= `vector.extract` $vector `` $position
33                     // attr-dict `:` type($vector) operation ::= `vector.load`
34                     // $base `[` $indices `]` attr-dict
35                     //               `:` type($base) `,` type($nresult)
36                     // operation ::= `vector.scalable.extract` $source `[` $pos
37                     // `]` attr-dict
38                     //               `:` type($res) `from` type($source)
39                     seq(choice('vector.extract', 'vector.load',
40                                'vector.scalable.extract'),
41                         field('operand', $.value_use),
42                         field('indices', $._dense_idx_list),
43                         field('attributes', optional($.attribute)),
44                         field('return', $._type_annotation)),
45
46                     // operation ::= `vector.fma` $lhs `,` $rhs `,` $acc
47                     // attr-dict `:` type($lhs)
48                     seq('vector.fma', field('lhs', $.value_use), ',',
49                         field('rhs', $.value_use), ',',
50                         field('acc', $.value_use),
51                         field('attributes', optional($.attribute)),
52                         field('return', $._type_annotation)),
53
54                     // operation ::= `vector.flat_transpose` $matrix attr-dict
55                     // `:` type($matrix) `->` type($res)
56                     seq('vector.flat_transpose', field('matrix', $.value_use),
57                         field('attributes', optional($.attribute)),
58                         field('return', $._function_type_annotation)),
59
60                     // operation ::= `vector.insert` $source `,` $dest
61                     // $position attr-dict
62                     //               `:` type($source) `into` type($dest)
63                     // operation ::= `vector.scalable.insert` $source `,` $dest
64                     // `[` $pos `]` attr-dict
65                     //               `:` type($source) `into` type($dest)
66                     // operation ::= `vector.shuffle` operands $mask attr-dict
67                     // `:` type(operands) operation ::= `vector.store`
68                     // $valueToStore `,` $base `[` $indices `]` attr-dict
69                     //               `:` type($base) `,` type($valueToStore)
70                     seq(choice('vector.insert', 'vector.scalable.insert',
71                                'vector.shuffle', 'vector.store'),
72                         field('source', $.value_use), ',',
73                         field('destination', $.value_use),
74                         field('position', $._dense_idx_list),
75                         field('attributes', optional($.attribute)),
76                         field('return', $._type_annotation)),
77
78                     // operation ::= `vector.insert_strided_slice` $source `,`
79                     // $dest attr-dict
80                     //               `:` type($source) `into` type($dest)
81                     seq('vector.insert_strided_slice',
82                         field('source', $.value_use), ',',
83                         field('destination', $.value_use),
84                         field('attributes', optional($.attribute)),
85                         field('return', $._type_annotation)),
86
87                     // operation ::= `vector.matrix_multiply` $lhs `,` $rhs
88                     // attr-dict
89                     //                `:` `(` type($lhs) `,` type($rhs) `)`
90                     //                `->` type($res)
91                     seq('vector.matrix_multiply', field('lhs', $.value_use),
92                         ',', field('rhs', $.value_use),
93                         field('attributes', optional($.attribute)),
94                         field('return', $._function_type_annotation)),
95
96                     // operation ::= `vector.print` $source attr-dict `:`
97                     // type($source)
98                     seq(choice('vector.print', 'vector.splat'),
99                         field('operand', $.value_use),
100                         field('attributes', optional($.attribute)),
101                         field('return', $._type_annotation)),
102
103                     seq('vector.transfer_read',
104                         field('source', seq($.value_use, $._dense_idx_list)),
105                         field('paddingMask', repeat(seq(',', $.value_use))),
106                         field('attributes', optional($.attribute)),
107                         field('return', $._type_annotation)),
108
109                     seq('vector.transfer_write', field('vector', $.value_use),
110                         ',',
111                         field('source', seq($.value_use, $._dense_idx_list)),
112                         field('mask', optional(seq(',', $.value_use))),
113                         field('attributes', optional($.attribute)),
114                         field('return', $._type_annotation)),
115
116                     // operation ::= `vector.yield` attr-dict ($operands^ `:`
117                     // type($operands))?
118                     seq('vector.yield',
119                         field('attributes', optional($.attribute)),
120                         field('results', optional($._value_use_type_list)))))
121}
122