mlir/utils/tree-sitter-mlir/dialect/linalg.js

   1 'use strict';
   2
   3 module.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 }