202 lines · javascript
1'use strict';2 3module.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 `,` $index19 // `:` 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` $src27 // $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-dict44 // `:` 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` $dest67 // ``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` $source76 // `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-dict95 // `:` 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-dict105 // `:` 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-dict118 // `:` 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` $source128 // (`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-dict144 // `:` 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` $source157 // (`padding_value` `(` $padding_value^ `:`158 // type($padding_value) `)`)?159 // (`outer_dims_perm` `=` $outer_dims_perm^)?160 // `inner_dims_pos` `=` $inner_dims_pos161 // `inner_tiles` `=`162 // custom<DynamicIndexList>($inner_tiles,163 // $static_inner_tiles) `into` $dest164 // attr-dict `:` type($source) `->`165 // type($dest)166 // operation ::= `tensor.unpack` $source167 // (`outer_dims_perm` `=` $outer_dims_perm^)?168 // `inner_dims_pos` `=` $inner_dims_pos169 // `inner_tiles` `=`170 // custom<DynamicIndexList>($inner_tiles,171 // $static_inner_tiles) `into` $dest172 // 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 $body187 // 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` $value196 // 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}202