xref: /llvm-project/mlir/utils/tree-sitter-mlir/dialect/affine.js (revision 10f8be19e7562624c0e63872e28cc963e03e56a5)
1'use strict';
2
3module.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}
119