xref: /llvm-project/mlir/utils/tree-sitter-mlir/dialect/linalg.js (revision 10f8be19e7562624c0e63872e28cc963e03e56a5)
1'use strict';
2
3module.exports = {
4  linalg_dialect : $ => prec.right(choice(
5                     seq(choice(
6                             'linalg.batch_matmul',
7                             'linalg.batch_matmul_transpose_b',
8                             'linalg.batch_matvec',
9                             'linalg.batch_reduce_matmul', 'linalg.broadcast',
10                             'linalg.conv_1d_ncw_fcw', 'linalg.conv_1d_nwc_wcf',
11                             'linalg.conv_1d', 'linalg.conv_2d_nchw_fchw',
12                             'linalg.conv_2d_ngchw_fgchw',
13                             'linalg.conv_2d_nhwc_fhwc',
14                             'linalg.conv_2d_nhwc_hwcf',
15                             'linalg.conv_2d_nhwc_hwcf_q', 'linalg.conv_2d',
16                             'linalg.conv_3d_ndhwc_dhwcf',
17                             'linalg.conv_3d_ndhwc_dhwcf_q', 'linalg.conv_3d',
18                             'linalg.copy', 'linalg.depthwise_conv_1d_nwc_wc',
19                             'linalg.depthwise_conv_1d_nwc_wcm',
20                             'linalg.depthwise_conv_2d_nchw_chw',
21                             'linalg.depthwise_conv_2d_nhwc_hwc',
22                             'linalg.depthwise_conv_2d_nhwc_hwc_q',
23                             'linalg.depthwise_conv_2d_nhwc_hwcm',
24                             'linalg.depthwise_conv_2d_nhwc_hwcm_q',
25                             'linalg.depthwise_conv_3d_ndhwc_dhwc',
26                             'linalg.depthwise_conv_3d_ndhwc_dhwcm',
27                             'linalg.dot', 'linalg.elemwise_binary',
28                             'linalg.elemwise_unary', 'linalg.fill',
29                             'linalg.fill_rng_2d', 'linalg.matmul',
30                             'linalg.matmul_transpose_b',
31                             'linalg.matmul_unsigned', 'linalg.matvec',
32                             'linalg.mmt4d', 'linalg.pooling_nchw_max',
33                             'linalg.pooling_nchw_sum',
34                             'linalg.pooling_ncw_max', 'linalg.pooling_ncw_sum',
35                             'linalg.pooling_ndhwc_max',
36                             'linalg.pooling_ndhwc_min',
37                             'linalg.pooling_ndhwc_sum',
38                             'linalg.pooling_nhwc_max',
39                             'linalg.pooling_nhwc_max_unsigned',
40                             'linalg.pooling_nhwc_min',
41                             'linalg.pooling_nhwc_min_unsigned',
42                             'linalg.pooling_nhwc_sum',
43                             'linalg.pooling_nwc_max',
44                             'linalg.pooling_nwc_max_unsigned',
45                             'linalg.pooling_nwc_min',
46                             'linalg.pooling_nwc_min_unsigned',
47                             'linalg.pooling_nwc_sum',
48                             'linalg.quantized_batch_matmul',
49                             'linalg.quantized_matmul', 'linalg.vecmat'),
50                         repeat1($._ins_outs_attributes),
51                         field('return', optional($._function_return))),
52
53                     seq('linalg.generic', repeat1($._ins_outs_attributes),
54                         field('body', $.region),
55                         field('return', optional($._function_return))),
56
57                     // operation ::= `linalg.index` $dim attr-dict `:`
58                     // type($result)
59                     seq('linalg.index', field('dimension', $.integer_literal),
60                         field('attributes', optional($.attribute)),
61                         field('return', $._type_annotation)),
62
63                     seq(choice('linalg.map', 'linalg.reduce'),
64                         repeat1($._ins_outs_attributes),
65                         field('arguments', $.block_arg_list),
66                         field('body', $.region),
67                         field('return', optional($._function_return))),
68
69                     seq('linalg.yield',
70                         field('attributes', optional($.attribute)),
71                         field('results', optional($._value_use_type_list))))),
72
73  _ins_outs_attributes : $ => choice($._ins, $._outs, $.attribute,
74                                     $._attribute_entry_list),
75  _ins : $ => seq(token('ins'), '(', $._value_use_type_list, ')'),
76  _outs : $ => seq(token('outs'), '(', $._value_use_type_list, ')'),
77  _attribute_entry_list : $ => seq($.bare_attribute_entry,
78                                   repeat(seq(',', $.bare_attribute_entry))),
79  bare_attribute_entry : $ => seq($.bare_id, '=', $.attribute_value)
80}
81