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

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