@felixtensor/tree-sitter-mlir 0.1.2 → 0.1.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/grammar.js CHANGED
@@ -2,16 +2,20 @@
2
2
  // @ts-check
3
3
 
4
4
  export default grammar({
5
- name: 'mlir',
6
- extras: $ => [/\s/, $.comment],
7
- conflicts: $ => [
5
+ name: "mlir",
6
+ extras: ($) => [/[\s\x00]/, $.comment],
7
+ conflicts: ($) => [
8
8
  [$._static_dim_list, $._static_dim_list],
9
9
  [$.dictionary_attribute, $.region],
10
10
  [$.custom_op_name, $.attribute_entry],
11
11
  [$.type_alias, $.dialect_namespace],
12
12
  [$.dialect_namespace, $.attribute_alias],
13
13
  [$.pretty_dialect_item],
14
- [$.array_literal, $._custom_body_element],
14
+ [$.array_literal, $._custom_body_element_base],
15
+ [$._custom_body_element_base, $.tensor_type],
16
+ [$._custom_body_element_base, $._custom_body_arrow],
17
+ [$._generic_custom_operation_with_location_attr_dict, $.custom_op_name],
18
+ [$._custom_body_dict_key, $.attribute_entry],
15
19
  [$._value_use_list, $._value_use_and_type],
16
20
  [$._type_list_no_parens, $._type_or_func_type],
17
21
  [$._type_list_parens, $._multi_dim_affine_expr_parens],
@@ -32,12 +36,28 @@ export default grammar({
32
36
  rules: {
33
37
  // =========================================================================
34
38
  // Top level production:
35
- // (operation | attribute-alias-def | type-alias-def)
39
+ // (operation | attribute-alias-def | type-alias-def)*
36
40
  // =========================================================================
37
- toplevel: $ => seq($._toplevel, repeat($._toplevel)),
38
- _toplevel: $ => choice($.operation, $.attribute_alias_def, $.type_alias_def,
39
- $.external_resources),
40
- external_resources: $ => seq('{-#', repeat($._pretty_dialect_item_contents), '#-}'),
41
+ toplevel: ($) => repeat($._toplevel),
42
+ _toplevel: ($) =>
43
+ choice(
44
+ $.operation,
45
+ $.attribute_alias_def,
46
+ $.type_alias_def,
47
+ $.external_resources,
48
+ ),
49
+ external_resources: ($) =>
50
+ seq("{-#", repeat($._external_resource_content), "#-}"),
51
+ _external_resource_content: ($) =>
52
+ choice(
53
+ $.string_literal,
54
+ // File metadata is YAML-like, not a dialect pretty body. Keep it loose
55
+ // so resource keys such as dense_resource_test_2xi32 are not split by
56
+ // dimension-list/type rules.
57
+ token(prec(-1, /[^#"\/]+/)),
58
+ "#",
59
+ "/",
60
+ ),
41
61
 
42
62
  // =========================================================================
43
63
  // Common syntax (lang-ref)
@@ -47,44 +67,95 @@ export default grammar({
47
67
  // float-literal ::= [-+]?[0-9]+[.][0-9]*([eE][-+]?[0-9]+)?
48
68
  // string-literal ::= `"` [^"\n\f\v\r]* `"`
49
69
  // =========================================================================
50
- _digit: $ => /[0-9]/,
51
- integer_literal: $ => choice($._decimal_literal, $._hexadecimal_literal),
52
- _decimal_literal: $ => token(seq(optional(/[-+]/), repeat1(/[0-9]/))),
53
- _unsigned_decimal_literal: $ => token(repeat1(/[0-9]/)),
54
- _unsigned_integer_literal: $ => choice($._unsigned_decimal_literal, $._hexadecimal_literal),
55
- _hexadecimal_literal: $ => token(seq('0x', repeat1(/[0-9a-fA-F]/))),
56
- float_literal: $ => token(seq(
57
- optional(/[-+]/), repeat1(/[0-9]/), '.', repeat(/[0-9]/),
58
- optional(seq(/[eE]/, optional(/[-+]/), repeat1(/[0-9]/))))),
59
- string_literal: $ => seq(
60
- '"',
61
- repeat(choice(/[^\\"\n\f\v\r]+/, $.escape_sequence, $.invalid_escape)),
62
- '"'
63
- ),
64
- escape_sequence: $ => token(seq('\\', choice(/[nt"\\]/, /[0-9a-fA-F]{2}/))),
65
- invalid_escape: $ => token(seq('\\', /[^\n\f\v\r]/)),
66
- bool_literal: $ => token(choice('true', 'false')),
67
- unit_literal: $ => token('unit'),
68
- uninitialized_literal: $ => token('uninitialized'),
69
- complex_literal: $ => seq('(', choice($.integer_literal, $.float_literal, $.bool_literal), ',',
70
- choice($.integer_literal, $.float_literal, $.bool_literal), ')'),
71
- tensor_literal: $ => seq(token(choice('dense', 'sparse')), '<',
72
- optional(seq(
73
- optional(seq($.type, ':')),
74
- seq($._tensor_literal_element, repeat(seq(',', $._tensor_literal_element)))
75
- )), '>'),
76
- _tensor_literal_element: $ => choice($.nested_idx_list, $._primitive_element),
77
- array_literal: $ => seq(token('array'), '<', $.type, optional(seq(':', $._idx_list)), '>'),
78
- _literal: $ => choice($.integer_literal, $.float_literal, $.string_literal, $.bool_literal,
79
- $.tensor_literal, $.array_literal, $.unit_literal, $.uninitialized_literal),
80
-
81
- nested_idx_list: $ => seq('[', optional(choice($.nested_idx_list, $._idx_list)),
82
- repeat(seq(',', $.nested_idx_list)), ']'),
83
- _idx_list: $ => prec.right(seq($._primitive_element,
84
- repeat(seq(',', $._primitive_element)))),
85
- _primitive_element: $ => seq($._primitive_idx_literal, optional(seq(':', $.type))),
86
- _primitive_idx_literal: $ => choice($.integer_literal, $.float_literal,
87
- $.bool_literal, $.complex_literal, $.string_literal),
70
+ _digit: ($) => /[0-9]/,
71
+ integer_literal: ($) => choice($._decimal_literal, $._hexadecimal_literal),
72
+ _decimal_literal: ($) => token(seq(optional(/[-+]/), repeat1(/[0-9]/))),
73
+ _unsigned_decimal_literal: ($) => token(repeat1(/[0-9]/)),
74
+ _unsigned_integer_literal: ($) =>
75
+ choice($._unsigned_decimal_literal, $._hexadecimal_literal),
76
+ _hexadecimal_literal: ($) => token(seq("0x", repeat1(/[0-9a-fA-F]/))),
77
+ float_literal: ($) =>
78
+ token(
79
+ seq(
80
+ optional(/[-+]/),
81
+ repeat1(/[0-9]/),
82
+ ".",
83
+ repeat(/[0-9]/),
84
+ optional(seq(/[eE]/, optional(/[-+]/), repeat1(/[0-9]/))),
85
+ ),
86
+ ),
87
+ string_literal: ($) =>
88
+ seq(
89
+ '"',
90
+ repeat(choice(/[^\\"\n\f\v\r]+/, $.escape_sequence, $.invalid_escape)),
91
+ '"',
92
+ ),
93
+ escape_sequence: ($) =>
94
+ token(seq("\\", choice(/[nt"\\]/, /[0-9a-fA-F]{2}/))),
95
+ invalid_escape: ($) => token(seq("\\", /[^\n\f\v\r]/)),
96
+ bool_literal: ($) => token(choice("true", "false")),
97
+ unit_literal: ($) => token("unit"),
98
+ uninitialized_literal: ($) => token("uninitialized"),
99
+ complex_literal: ($) =>
100
+ seq(
101
+ "(",
102
+ choice($.integer_literal, $.float_literal, $.bool_literal),
103
+ ",",
104
+ choice($.integer_literal, $.float_literal, $.bool_literal),
105
+ ")",
106
+ ),
107
+ tensor_literal: ($) =>
108
+ seq(
109
+ choice($._dense_keyword, $._sparse_keyword),
110
+ "<",
111
+ optional(
112
+ seq(
113
+ optional(seq($.type, ":")),
114
+ seq(
115
+ $._tensor_literal_element,
116
+ repeat(seq(",", $._tensor_literal_element)),
117
+ ),
118
+ ),
119
+ ),
120
+ ">",
121
+ ),
122
+ _tensor_literal_element: ($) =>
123
+ choice($.nested_idx_list, $._primitive_element),
124
+ array_literal: ($) =>
125
+ seq(token("array"), "<", $.type, optional(seq(":", $._idx_list)), ">"),
126
+ _literal: ($) =>
127
+ choice(
128
+ $.integer_literal,
129
+ $.float_literal,
130
+ $.string_literal,
131
+ $.bool_literal,
132
+ $.tensor_literal,
133
+ $.array_literal,
134
+ $.unit_literal,
135
+ $.uninitialized_literal,
136
+ ),
137
+
138
+ nested_idx_list: ($) =>
139
+ seq(
140
+ "[",
141
+ optional(choice($.nested_idx_list, $._idx_list)),
142
+ repeat(seq(",", $.nested_idx_list)),
143
+ "]",
144
+ ),
145
+ _idx_list: ($) =>
146
+ prec.right(
147
+ seq($._primitive_element, repeat(seq(",", $._primitive_element))),
148
+ ),
149
+ _primitive_element: ($) =>
150
+ seq($._primitive_idx_literal, optional(seq(":", $.type))),
151
+ _primitive_idx_literal: ($) =>
152
+ choice(
153
+ $.integer_literal,
154
+ $.float_literal,
155
+ $.bool_literal,
156
+ $.complex_literal,
157
+ $.string_literal,
158
+ ),
88
159
 
89
160
  // =========================================================================
90
161
  // Identifiers
@@ -93,16 +164,28 @@ export default grammar({
93
164
  // suffix-id ::= (digit+ | ((letter|id-punct) (letter|id-punct|digit)*))
94
165
  // symbol-ref-id ::= `@` (suffix-id | string-literal) (`::` symbol-ref-id)?
95
166
  // =========================================================================
96
- bare_id: $ => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.]/))),
97
- _alias_or_dialect_id: $ => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/))),
98
- bare_id_list: $ => seq($.bare_id, repeat(seq(',', $.bare_id))),
99
- value_use: $ => seq('%', $._suffix_id),
100
- _suffix_id: $ => token(seq(choice(repeat1(/[0-9]/),
101
- seq(/[a-zA-Z_$.-]/, repeat(/[a-zA-Z0-9_$.-]/))),
102
- optional(seq(choice(':', '#'), repeat1(/[0-9]/))))),
103
- symbol_ref_id: $ => seq('@', choice($._suffix_id, $.string_literal),
104
- optional(seq('::', $.symbol_ref_id))),
105
- _value_use_list: $ => seq($.value_use, repeat(seq(',', $.value_use))),
167
+ bare_id: ($) => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.]/))),
168
+ _alias_or_dialect_id: ($) =>
169
+ token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/))),
170
+ bare_id_list: ($) => seq($.bare_id, repeat(seq(",", $.bare_id))),
171
+ value_use: ($) => seq("%", $._suffix_id),
172
+ _suffix_id: ($) =>
173
+ token(
174
+ seq(
175
+ choice(
176
+ repeat1(/[0-9]/),
177
+ seq(/[a-zA-Z_$.-]/, repeat(/[a-zA-Z0-9_$.-]/)),
178
+ ),
179
+ optional(seq(choice(":", "#"), repeat1(/[0-9]/))),
180
+ ),
181
+ ),
182
+ symbol_ref_id: ($) =>
183
+ seq(
184
+ "@",
185
+ choice($._suffix_id, $.string_literal),
186
+ optional(seq("::", $.symbol_ref_id)),
187
+ ),
188
+ _value_use_list: ($) => seq($.value_use, repeat(seq(",", $.value_use))),
106
189
 
107
190
  // =========================================================================
108
191
  // Operations
@@ -113,58 +196,86 @@ export default grammar({
113
196
  // dictionary-attribute? `:` function-type
114
197
  // custom-operation ::= bare-id custom-operation-format
115
198
  // =========================================================================
116
- operation: $ => seq(
117
- field('lhs', optional($._op_result_list)),
118
- field('rhs', choice($.generic_operation, $.custom_operation)),
119
- field('location', optional($.trailing_location))),
120
-
121
- generic_operation: $ =>
122
- seq($.string_literal, $._value_use_list_parens, optional($._successor_list),
123
- optional($.properties), optional($._region_list), optional($.attribute),
124
- ':', $.function_type),
125
-
126
- _op_result_list: $ => seq($.op_result, repeat(seq(',', $.op_result)), '='),
127
- op_result: $ => seq($.value_use, optional(seq(':', $.integer_literal))),
128
- _successor_list: $ => seq('[', $.successor, repeat(seq(',', $.successor)), ']'),
129
- successor: $ => prec.right(seq($.caret_id, optional($._value_arg_list))),
130
- _region_list: $ => seq('(', $.region, repeat(seq(',', $.region)), ')'),
131
- properties: $ => seq('<{', optional($.attribute_entry),
132
- repeat(seq(',', $.attribute_entry)), '}>'),
133
- dictionary_attribute: $ => seq('{', optional($.attribute_entry),
134
- repeat(seq(',', $.attribute_entry)), '}'),
135
- trailing_location: $ => seq(token('loc'), '(', $.location, ')'),
199
+ operation: ($) =>
200
+ seq(
201
+ field("lhs", optional($._op_result_list)),
202
+ field("rhs", choice($.generic_operation, $.custom_operation)),
203
+ field("location", optional($.trailing_location)),
204
+ ),
205
+
206
+ generic_operation: ($) =>
207
+ seq(
208
+ $.string_literal,
209
+ $._value_use_list_parens,
210
+ optional($._successor_list),
211
+ optional($.properties),
212
+ optional($._region_list),
213
+ optional($.attribute),
214
+ ":",
215
+ $.function_type,
216
+ ),
217
+
218
+ _op_result_list: ($) =>
219
+ seq($.op_result, repeat(seq(",", $.op_result)), "="),
220
+ op_result: ($) => seq($.value_use, optional(seq(":", $.integer_literal))),
221
+ _successor_list: ($) =>
222
+ seq("[", $.successor, repeat(seq(",", $.successor)), "]"),
223
+ successor: ($) => prec.right(seq($.caret_id, optional($._value_arg_list))),
224
+ _region_list: ($) => seq("(", $.region, repeat(seq(",", $.region)), ")"),
225
+ properties: ($) =>
226
+ seq(
227
+ "<{",
228
+ optional($.attribute_entry),
229
+ repeat(seq(",", $.attribute_entry)),
230
+ "}>",
231
+ ),
232
+ dictionary_attribute: ($) =>
233
+ seq(
234
+ "{",
235
+ optional($.attribute_entry),
236
+ repeat(seq(",", $.attribute_entry)),
237
+ "}",
238
+ ),
239
+ trailing_location: ($) => seq(token("loc"), "(", $.location, ")"),
136
240
  // LangRef "Source Locations": location ::= unknown | name | filelinecol |
137
241
  // filelocrange | callsite | fused | loc-alias-ref
138
- location: $ => choice(
139
- $.unknown_location,
140
- $._string_location,
141
- $.callsite_location,
142
- $.fused_location,
143
- $.attribute_alias,
144
- $.dialect_attribute,
145
- ),
146
- unknown_location: $ => token('unknown'),
147
- callsite_location: $ => seq(token('callsite'), '(',
148
- $.location, token('at'), $.location, ')'),
149
- fused_location: $ => seq(token('fused'),
150
- optional(seq('<', $.attribute_value, '>')),
151
- '[', $.location, repeat(seq(',', $.location)), ']'),
242
+ location: ($) =>
243
+ choice(
244
+ $.unknown_location,
245
+ $._string_location,
246
+ $.callsite_location,
247
+ $.fused_location,
248
+ $.attribute_alias,
249
+ $.dialect_attribute,
250
+ ),
251
+ unknown_location: ($) => token("unknown"),
252
+ callsite_location: ($) =>
253
+ seq(token("callsite"), "(", $.location, token("at"), $.location, ")"),
254
+ fused_location: ($) =>
255
+ seq(
256
+ token("fused"),
257
+ optional(seq("<", $.attribute_value, ">")),
258
+ "[",
259
+ $.location,
260
+ repeat(seq(",", $.location)),
261
+ "]",
262
+ ),
152
263
  // Combines name-location and filelinecol-location (LangRef uses the same
153
264
  // string-literal prefix; what follows distinguishes them).
154
- _string_location: $ => seq(
155
- $.string_literal,
156
- optional(choice(
157
- $._filelinecol_suffix,
158
- seq('(', $.location, ')'),
159
- )),
160
- ),
161
- _filelinecol_suffix: $ => seq(
162
- ':', $.integer_literal,
163
- optional(seq(':', $.integer_literal)),
164
- optional(seq(token('to'),
165
- optional($.integer_literal),
166
- ':', $.integer_literal)),
167
- ),
265
+ _string_location: ($) =>
266
+ seq(
267
+ $.string_literal,
268
+ optional(choice($._filelinecol_suffix, seq("(", $.location, ")"))),
269
+ ),
270
+ _filelinecol_suffix: ($) =>
271
+ seq(
272
+ ":",
273
+ $.integer_literal,
274
+ optional(seq(":", $.integer_literal)),
275
+ optional(
276
+ seq(token("to"), optional($.integer_literal), ":", $.integer_literal),
277
+ ),
278
+ ),
168
279
 
169
280
  // =========================================================================
170
281
  // Three-tier custom operation system
@@ -172,41 +283,67 @@ export default grammar({
172
283
  // Tier 1: func_operation, module_operation — structural ops for navigation
173
284
  // Tier 2: _generic_custom_operation — all other dialect.op_name patterns
174
285
  // =========================================================================
175
- custom_operation: $ => choice(
176
- prec(2, $.func_operation),
177
- prec(2, $.module_operation),
178
- prec(2, $._affine_for_operation),
179
- $._generic_custom_operation,
180
- ),
286
+ custom_operation: ($) =>
287
+ choice(
288
+ prec(2, $.func_operation),
289
+ prec(2, $.module_operation),
290
+ prec(2, $._affine_for_operation),
291
+ $._pdl_interp_record_match_operation,
292
+ $._generic_custom_operation_with_location_attr_dict,
293
+ $._generic_custom_operation,
294
+ ),
181
295
 
182
296
  // Tier 1: Function operations (func.func, llvm.func)
183
297
  // Token prec 20 ensures these keywords win over _dotted_op_name (prec 10)
184
- func_operation: $ => prec.right(seq(
185
- field('name', choice(token(prec(20, 'func.func')), token(prec(20, 'llvm.func')))),
186
- field('visibility', optional(choice('private', 'public'))),
187
- // Optional leading specifier keywords between the function name and its
188
- // symbol. Covers MLIR symbol visibility (`nested`) and the LLVM dialect's
189
- // llvm.func linkage / calling-convention / unnamed_addr / visibility
190
- // keywords (internal, external, linkonce, fastcc, amdgpu_kernelcc,
191
- // local_unnamed_addr, hidden, ...). MLIR's CConv enum alone has ~50
192
- // keywords, so rather than enumerate a brittle list we accept bare_id-
193
- // shaped specifiers here; the trailing `@symbol` (symbol_ref_id) is an
194
- // unambiguous terminator.
195
- repeat(field('specifier', alias($.bare_id, $.function_specifier))),
196
- field('sym_name', $.symbol_ref_id),
197
- field('arguments', $.func_arg_list),
198
- field('return', optional($.func_return)),
199
- field('attributes', optional(seq(optional(token('attributes')), $.dictionary_attribute))),
200
- field('body', optional($.region)),
201
- )),
298
+ func_operation: ($) =>
299
+ prec.right(
300
+ seq(
301
+ field(
302
+ "name",
303
+ choice(token(prec(20, "func.func")), token(prec(20, "llvm.func"))),
304
+ ),
305
+ field("visibility", optional(choice("private", "public"))),
306
+ // Optional leading specifier keywords between the function name and its
307
+ // symbol. Covers MLIR symbol visibility (`nested`) and the LLVM dialect's
308
+ // llvm.func linkage / calling-convention / unnamed_addr / visibility
309
+ // keywords (internal, external, linkonce, fastcc, amdgpu_kernelcc,
310
+ // local_unnamed_addr, hidden, ...). MLIR's CConv enum alone has ~50
311
+ // keywords, so rather than enumerate a brittle list we accept bare_id-
312
+ // shaped specifiers here; the trailing `@symbol` (symbol_ref_id) is an
313
+ // unambiguous terminator.
314
+ repeat(field("specifier", alias($.bare_id, $.function_specifier))),
315
+ field("sym_name", $.symbol_ref_id),
316
+ field("arguments", $.func_arg_list),
317
+ field("return", optional($.func_return)),
318
+ field(
319
+ "attributes",
320
+ optional(
321
+ seq(optional(token("attributes")), $.dictionary_attribute),
322
+ ),
323
+ ),
324
+ field("body", optional($.region)),
325
+ ),
326
+ ),
202
327
 
203
328
  // Tier 1: Module operations (module, builtin.module)
204
- module_operation: $ => prec.right(seq(
205
- field('name', choice(token(prec(20, 'builtin.module')), token(prec(20, 'module')))),
206
- field('sym_name', optional($.symbol_ref_id)),
207
- field('attributes', optional(seq(optional(token('attributes')), $.attribute))),
208
- field('body', $.region),
209
- )),
329
+ module_operation: ($) =>
330
+ prec.right(
331
+ seq(
332
+ field(
333
+ "name",
334
+ choice(
335
+ token(prec(20, "builtin.module")),
336
+ token(prec(20, "module")),
337
+ ),
338
+ ),
339
+ field("sym_name", optional($.symbol_ref_id)),
340
+ field(
341
+ "attributes",
342
+ optional(seq(optional(token("attributes")), $.attribute)),
343
+ ),
344
+ field("body", $.region),
345
+ ),
346
+ ),
210
347
 
211
348
  // affine.for prints its induction-variable location before the body
212
349
  // (`affine.for %i loc("iv") = 0 to 8 { ... } loc(...)`). The generic
@@ -214,21 +351,58 @@ export default grammar({
214
351
  // operation; without this structured rule, that machinery would bind the
215
352
  // induction loc to the operation slot and leave the operation's own
216
353
  // trailing loc unparseable.
217
- _affine_for_operation: $ => prec.right(seq(
218
- field('name', alias(token(prec(20, 'affine.for')), $.custom_op_name)),
219
- $.value_use,
220
- optional(field('induction_location', $.trailing_location)),
221
- repeat($._custom_body_element),
222
- )),
354
+ _affine_for_operation: ($) =>
355
+ prec.right(
356
+ seq(
357
+ field("name", alias(token(prec(20, "affine.for")), $.custom_op_name)),
358
+ $.value_use,
359
+ optional(field("induction_location", $.trailing_location)),
360
+ repeat($._custom_body_element),
361
+ ),
362
+ ),
363
+
364
+ _pdl_interp_record_match_operation: ($) =>
365
+ prec.dynamic(
366
+ -1,
367
+ prec.right(
368
+ seq(
369
+ field(
370
+ "name",
371
+ alias(
372
+ token(prec(20, "pdl_interp.record_match")),
373
+ $.custom_op_name,
374
+ ),
375
+ ),
376
+ repeat($._custom_body_element),
377
+ $._custom_body_location_list,
378
+ repeat($._custom_body_element),
379
+ ),
380
+ ),
381
+ ),
223
382
 
224
383
  // Tier 2: Generic custom operation — dialect.op_name + structural body
225
384
  // Negative dynamic precedence makes the parser prefer ending the body
226
385
  // and starting a new operation (with _op_result_list) over extending
227
386
  // the body with more elements, when both paths are valid (GLR).
228
- _generic_custom_operation: $ => prec.dynamic(-1, prec.right(seq(
229
- field('name', $.custom_op_name),
230
- repeat($._custom_body_element),
231
- ))),
387
+ _generic_custom_operation: ($) =>
388
+ prec.dynamic(
389
+ -1,
390
+ prec.right(
391
+ seq(field("name", $.custom_op_name), repeat($._custom_body_element)),
392
+ ),
393
+ ),
394
+
395
+ _generic_custom_operation_with_location_attr_dict: ($) =>
396
+ prec.dynamic(
397
+ -1,
398
+ prec.right(
399
+ seq(
400
+ field("name", alias($._dotted_op_name, $.custom_op_name)),
401
+ $._custom_body_location_attr_dict,
402
+ repeat($._custom_body_element),
403
+ ),
404
+ ),
405
+ ),
232
406
 
233
407
  // Operation name: dotted form (arith.addi, scf.forall.in_parallel)
234
408
  // or known bare names (return, call, constant, etc.)
@@ -241,48 +415,160 @@ export default grammar({
241
415
  // prec.dynamic(-1) on _generic_custom_operation keeps bare_id as a body
242
416
  // element when inside a custom op body; it only acts as an op name at
243
417
  // region/block boundaries where operation+ is required.
244
- custom_op_name: $ => choice($._dotted_op_name, $._bare_op_name, $.bare_id),
245
- _dotted_op_name: $ => token(prec(10, seq(
246
- /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/),
247
- repeat1(seq('.', /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.]/))),
248
- ))),
418
+ custom_op_name: ($) =>
419
+ choice($._dotted_op_name, $._bare_op_name, $.bare_id),
420
+ _dotted_op_name: ($) =>
421
+ token(
422
+ prec(
423
+ 10,
424
+ seq(
425
+ /[a-zA-Z_]/,
426
+ repeat(/[a-zA-Z0-9_$]/),
427
+ repeat1(seq(".", /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.]/))),
428
+ ),
429
+ ),
430
+ ),
249
431
  // Bare operation names: spec-sanctioned no-prefix aliases only.
250
432
  // All other dialect ops go through _dotted_op_name or _generic_custom_operation.
251
- _bare_op_name: $ => token(prec(10, choice(
252
- 'return', 'call', 'call_indirect', 'constant', 'unrealized_conversion_cast',
253
- ))),
433
+ _bare_op_name: ($) =>
434
+ token(
435
+ prec(
436
+ 10,
437
+ choice(
438
+ "return",
439
+ "call",
440
+ "call_indirect",
441
+ "constant",
442
+ "unrealized_conversion_cast",
443
+ ),
444
+ ),
445
+ ),
254
446
 
255
447
  // Structural body elements that can appear in custom operation format.
256
448
  // These are recognized by sigils (%,^,@,!,#) or by structural delimiters.
257
- _custom_body_element: $ => choice(
258
- $.value_use, // %foo, %0
259
- $.symbol_ref_id, // @sym, @"string"
260
- $.successor, // ^bb0, ^bb0(%arg : type)
261
- prec(2, $.type), // !type, i32, memref<...>, etc.
262
- $.attribute, // #attr, {dict}, affine_map<...>
263
- $.region, // { ... } (regions with operations)
264
- $._custom_body_value_group, // {%v : type, ...}
265
- $._custom_body_paren, // ( ... )
266
- $._custom_body_bracket, // [ ... ]
267
- $._custom_body_angle_group, // < ... >
268
- $._literal, // 42, 3.14, "string", true, dense<...>
269
- 'array', // property names may collide with array<...>
270
- 'vector', // OpenACC keyword may collide with vector<...>
271
- 'ceildiv', 'floordiv', 'mod', // inline affine keywords
272
- $.bare_id, // keywords: to, from, step, ins, outs, etc.
273
- ',', '=', ':', '->', '*', '?', $.dimension_separator, '+', '-', '/', '&', '|', '~',
274
- ),
275
-
276
- _custom_body_paren: $ => seq('(', repeat($._nested_custom_body_element), ')'),
277
- _custom_body_bracket: $ => seq('[', repeat($._nested_custom_body_element), ']'),
278
- _custom_body_value_group: $ => seq('{', $.value_use, ':', $.type,
279
- repeat(seq(',', $.value_use, ':', $.type)), '}'),
280
- _custom_body_angle_group: $ => seq('<', repeat($._nested_custom_body_element), '>'),
449
+ _custom_body_element: ($) =>
450
+ choice(
451
+ $._custom_body_element_base,
452
+ $._custom_body_successor_marker, // >^bb1 (WasmSSA if continuation)
453
+ ),
454
+
455
+ _custom_body_element_base: ($) =>
456
+ choice(
457
+ $.value_use, // %foo, %0
458
+ $.symbol_ref_id, // @sym, @"string"
459
+ $.successor, // ^bb0, ^bb0(%arg : type)
460
+ $._custom_body_complex_label, // complex: %value (IRDL operand label)
461
+ prec(2, $.type), // !type, i32, memref<...>, etc.
462
+ $.attribute, // #attr, {dict}, affine_map<...>
463
+ $._custom_body_tuple_group, // {(%v), (%w)}
464
+ $.region, // { ... } (regions with operations)
465
+ $._custom_body_value_group, // {%v : type, ...}
466
+ $._custom_body_ssa_dict, // {"attr" = %value, ...} / options with SSA values
467
+ $._custom_body_module_symbol_arg, // module(@sym) kernel attr
468
+ $._custom_body_sparse_operand, // sparse(%idx : type)
469
+ $._custom_body_paren, // ( ... )
470
+ $._custom_body_bracket, // [ ... ]
471
+ $._custom_body_angle_group, // < ... >
472
+ $.variadic, // custom assembly ellipsis marker
473
+ $._literal, // 42, 3.14, "string", true, dense<...>
474
+ "array", // property names may collide with array<...>
475
+ "vector", // OpenACC keyword may collide with vector<...>
476
+ "tensor", // AMDGPU/NVGPU keyword may collide with tensor<...>
477
+ "ceildiv",
478
+ "floordiv",
479
+ "mod", // inline affine keywords
480
+ $.bare_id, // keywords: to, from, step, ins, outs, etc.
481
+ $._custom_body_arrow, // <- (mapped-from, e.g. omp.fuse <- (...))
482
+ ",",
483
+ "=",
484
+ ":",
485
+ "->",
486
+ "*",
487
+ "?",
488
+ $.dimension_separator,
489
+ "+",
490
+ "-",
491
+ "/",
492
+ "&",
493
+ "|",
494
+ "~",
495
+ ),
496
+
497
+ _custom_body_paren: ($) =>
498
+ seq("(", repeat($._nested_custom_body_element), ")"),
499
+ _custom_body_bracket: ($) =>
500
+ seq("[", repeat($._nested_custom_body_element), "]"),
501
+ _custom_body_value_group: ($) =>
502
+ seq(
503
+ "{",
504
+ $.value_use,
505
+ ":",
506
+ $.type,
507
+ repeat(seq(",", $.value_use, ":", $.type)),
508
+ "}",
509
+ ),
510
+ _custom_body_tuple_group: ($) =>
511
+ seq(
512
+ "{",
513
+ $._custom_body_tuple,
514
+ repeat(seq(",", $._custom_body_tuple)),
515
+ "}",
516
+ ),
517
+ _custom_body_tuple: ($) => seq("(", $._value_use_list, ")"),
518
+ _custom_body_location_attr_dict: ($) =>
519
+ seq($.trailing_location, $.dictionary_attribute),
520
+ _custom_body_complex_label: ($) =>
521
+ prec(1, seq($._complex_label_start, $.value_use)),
522
+ _custom_body_ssa_dict: ($) =>
523
+ seq(
524
+ "{",
525
+ repeat(seq($._custom_body_attr_dict_entry, ",")),
526
+ $._custom_body_ssa_dict_entry,
527
+ repeat(seq(",", $._custom_body_mixed_dict_entry)),
528
+ "}",
529
+ ),
530
+ _custom_body_dict_key: ($) => $.string_literal,
531
+ _custom_body_attr_dict_entry: ($) =>
532
+ seq($._custom_body_dict_key, "=", $.attribute_value),
533
+ _custom_body_ssa_dict_entry: ($) =>
534
+ seq(
535
+ $._custom_body_dict_key,
536
+ "=",
537
+ choice($.value_use, $._custom_body_ssa_value_array),
538
+ ),
539
+ _custom_body_mixed_dict_entry: ($) =>
540
+ seq(
541
+ $._custom_body_dict_key,
542
+ "=",
543
+ choice($.attribute_value, $.value_use, $._custom_body_ssa_value_array),
544
+ ),
545
+ _custom_body_ssa_value_array: ($) =>
546
+ seq("[", $.value_use, repeat(seq(",", $.value_use)), "]"),
547
+ _custom_body_successor_marker: ($) => seq(">", $.successor),
548
+ _custom_body_module_symbol_arg: ($) =>
549
+ seq(token(prec(20, "module(")), $.symbol_ref_id, ")"),
550
+ _custom_body_sparse_operand: ($) =>
551
+ seq($._sparse_keyword, $._custom_body_paren),
552
+ _custom_body_location_list: ($) =>
553
+ seq(token("loc"), "(", "[", optional($._value_use_list), "]", ")"),
554
+ _custom_body_angle_group: ($) =>
555
+ seq("<", repeat($._nested_custom_body_element), ">"),
556
+ // "Mapped-from" arrow used by OpenMP loop-transform ops, e.g.
557
+ // omp.fuse(%fused) <- (%loop0, %loop1)
558
+ // omp.tile(%grid, %intratile) <- (%loop) sizes(%ts : i32)
559
+ // Deliberately built from the existing single-char '<' and '-' tokens
560
+ // rather than a combined '<-' token: a length-2 '<-' token would win
561
+ // tree-sitter's longest-match at a dialect-attribute body boundary and
562
+ // mis-lex negative payloads like `#smt.bv<-1>` as `<-` `1`. Keeping the
563
+ // tokens separate leaves the lexer unchanged; GLR distinguishes the arrow
564
+ // from `_custom_body_angle_group` (which requires a closing '>').
565
+ _custom_body_arrow: ($) => seq("<", "-"),
281
566
  // Only nested groups accept `trailing_location` as a body element.
282
567
  // At the top level it is omitted on purpose so the operation rule
283
568
  // captures a trailing `loc(...)` as the operation's location instead of
284
569
  // it being swallowed by the body repetition.
285
- _nested_custom_body_element: $ => choice($._custom_body_element, $.trailing_location),
570
+ _nested_custom_body_element: ($) =>
571
+ choice($._custom_body_element_base, $.trailing_location),
286
572
 
287
573
  // =========================================================================
288
574
  // Blocks
@@ -290,241 +576,532 @@ export default grammar({
290
576
  // block-label ::= block-id block-arg-list? `:`
291
577
  // caret-id ::= `^` suffix-id
292
578
  // =========================================================================
293
- block: $ => seq($.block_label, repeat($.operation)),
294
- block_label: $ => seq($._block_id, optional($.block_arg_list), ':'),
295
- _block_id: $ => $.caret_id,
296
- caret_id: $ => seq('^', $._suffix_id),
297
- _value_use_and_type: $ => seq($.value_use, optional(seq(':', $.type)),
298
- optional($.trailing_location)),
299
- _value_use_and_type_list: $ => seq($._value_use_and_type,
300
- repeat(seq(',', $._value_use_and_type))),
301
- block_arg_list: $ => seq('(', optional($._value_use_and_type_list), ')'),
302
- _value_arg_list: $ => seq('(', optional(choice(
303
- $._value_use_type_list, // bulk format: (%v0, %v1 : t0, t1)
304
- $._value_use_and_type_list, // per-pair format: (%v0 : t0, %v1 : t1)
305
- )), ')'),
306
- _value_use_type_list: $ => seq($._value_use_list, ':', $._type_list_no_parens),
579
+ block: ($) => seq($.block_label, repeat($.operation)),
580
+ block_label: ($) => seq($._block_id, optional($.block_arg_list), ":"),
581
+ _block_id: ($) => $.caret_id,
582
+ caret_id: ($) => seq("^", $._suffix_id),
583
+ _value_use_and_type: ($) =>
584
+ seq(
585
+ $.value_use,
586
+ optional(seq(":", $.type)),
587
+ optional($.trailing_location),
588
+ ),
589
+ _value_use_and_type_list: ($) =>
590
+ seq($._value_use_and_type, repeat(seq(",", $._value_use_and_type))),
591
+ block_arg_list: ($) => seq("(", optional($._value_use_and_type_list), ")"),
592
+ _value_arg_list: ($) =>
593
+ seq(
594
+ "(",
595
+ optional(
596
+ choice(
597
+ $._value_use_type_list, // bulk format: (%v0, %v1 : t0, t1)
598
+ $._value_use_and_type_list, // per-pair format: (%v0 : t0, %v1 : t1)
599
+ ),
600
+ ),
601
+ ")",
602
+ ),
603
+ _value_use_type_list: ($) =>
604
+ seq($._value_use_list, ":", $._type_list_no_parens),
307
605
 
308
606
  // =========================================================================
309
607
  // Regions
310
608
  // region ::= `{` entry-block? block* `}`
311
609
  // entry-block ::= operation+
312
610
  // =========================================================================
313
- region: $ => seq('{', optional($.entry_block), repeat($.block), '}'),
314
- entry_block: $ => repeat1($.operation),
611
+ region: ($) => seq("{", optional($.entry_block), repeat($.block), "}"),
612
+ entry_block: ($) => repeat1($.operation),
315
613
 
316
614
  // =========================================================================
317
615
  // Types
318
616
  // type ::= type-alias | dialect-type | builtin-type
319
617
  // function-type ::= (type | type-list-parens) `->` (type | type-list-parens)
320
618
  // =========================================================================
321
- type: $ => choice($.type_alias, $.dialect_type, $.builtin_type),
322
- _type_list_no_parens: $ => prec.left(seq(choice($.type, $.function_type),
323
- repeat(seq(',', choice($.type, $.function_type))))),
324
- _type_list_parens: $ => seq('(', optional($._type_list_no_parens), ')'),
325
- function_type: $ => seq(choice($.type, $._type_list_parens), $._function_return),
326
- _function_return: $ => seq(token('->'), choice($.type, $._type_list_parens)),
327
- _type_annotation: $ => seq(':', $._type_list_no_parens),
328
- _function_type_annotation: $ => seq(':', $.function_type),
329
- _literal_and_type: $ => seq($._literal, optional($._type_annotation)),
619
+ type: ($) => choice($.type_alias, $.dialect_type, $.builtin_type),
620
+ _type_list_no_parens: ($) =>
621
+ prec.left(
622
+ seq(
623
+ choice($.type, $.function_type),
624
+ repeat(seq(",", choice($.type, $.function_type))),
625
+ ),
626
+ ),
627
+ _type_list_parens: ($) => seq("(", optional($._type_list_no_parens), ")"),
628
+ function_type: ($) =>
629
+ seq(choice($.type, $._type_list_parens), $._function_return),
630
+ _function_return: ($) =>
631
+ seq(token("->"), choice($.type, $._type_list_parens)),
632
+ _type_annotation: ($) => seq(":", $._type_list_no_parens),
633
+ _function_type_annotation: ($) => seq(":", $.function_type),
634
+ _literal_and_type: ($) => seq($._literal, optional($._type_annotation)),
330
635
 
331
636
  // Type aliases
332
- type_alias_def: $ => seq('!', $._alias_or_dialect_id, '=', $.type),
333
- type_alias: $ => seq('!', $._alias_or_dialect_id),
637
+ type_alias_def: ($) => seq("!", $._alias_or_dialect_id, "=", $.type),
638
+ type_alias: ($) => seq("!", $._alias_or_dialect_id),
334
639
 
335
640
  // Dialect Types
336
- dialect_type: $ => seq(
337
- '!', choice($.opaque_dialect_item, $.pretty_dialect_item, $.parametric_dialect_item)),
338
- dialect_namespace: $ => $._alias_or_dialect_id,
339
- dialect_ident: $ => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.-]/))),
340
- opaque_dialect_item: $ => prec(1, seq($.dialect_namespace, '<', $.string_literal, '>')),
341
- pretty_dialect_item: $ => seq($.dialect_namespace, '.', $.dialect_ident,
342
- optional($.pretty_dialect_item_body)),
641
+ dialect_type: ($) =>
642
+ seq(
643
+ "!",
644
+ choice(
645
+ $.opaque_dialect_item,
646
+ $.pretty_dialect_item,
647
+ $.parametric_dialect_item,
648
+ ),
649
+ ),
650
+ dialect_namespace: ($) => $._alias_or_dialect_id,
651
+ dialect_ident: ($) => token(seq(/[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.-]/))),
652
+ opaque_dialect_item: ($) =>
653
+ prec(1, seq($.dialect_namespace, "<", $.string_literal, ">")),
654
+ pretty_dialect_item: ($) =>
655
+ seq(
656
+ $.dialect_namespace,
657
+ ".",
658
+ $.dialect_ident,
659
+ optional($.pretty_dialect_item_body),
660
+ ),
343
661
  // Parametric dialect item: !namespace<...> or #namespace<...> without dot-separated ident
344
- parametric_dialect_item: $ => seq($.dialect_namespace, $.pretty_dialect_item_body),
345
- pretty_dialect_item_body: $ => seq('<', repeat($._pretty_dialect_item_contents), '>'),
346
- _pretty_dialect_item_contents: $ => prec.left(choice(
347
- $.pretty_dialect_item_body,
348
- $._pretty_dialect_bang_body_token,
349
- $._pretty_dialect_body_attribute,
350
- $.dialect_dim_list,
351
- $.type,
352
- prec(2, $.attribute),
353
- $._literal,
354
- 'dense', 'sparse', 'array', 'vector', 'tensor', 'opaque',
355
- $.bare_id,
356
- ',', ':', '=', '->', '(', ')', '[', ']', '{', '}', '*', '?',
357
- '@', '#',
358
- token(prec(-1, /[^<>]/))
359
- )),
360
- _pretty_dialect_bang_body_token: $ => token(prec(1,
361
- seq('!', /[^a-zA-Z_<>]/))),
362
- _pretty_dialect_body_attribute: $ => choice(
363
- alias($._pretty_dialect_attribute_token, $.dialect_attribute),
364
- alias($._pretty_attribute_alias_token, $.attribute_alias)),
365
- _pretty_attribute_alias_token: $ => token(prec(2,
366
- seq('#', /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/)))),
367
- _pretty_dialect_attribute_token: $ => token(prec(2,
368
- seq('#', /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/), '.',
369
- /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$.-]/)))),
370
- dimension_separator: $ => token(prec(10, 'x')),
371
- dialect_dim_list: $ => prec.dynamic(1, seq($._dialect_dim_primitive,
372
- repeat1(seq($.dimension_separator, $._dialect_dim_primitive)))),
373
- _dialect_dim_primitive: $ => choice(
374
- prec(1, $.type),
375
- prec(1, $._pretty_dialect_body_type),
376
- alias($.integer_literal, $.dimension_size),
377
- '?',
378
- '*'),
379
- _pretty_dialect_body_type: $ => seq(
380
- choice($.bare_id, 'array'),
381
- $.pretty_dialect_item_body),
662
+ parametric_dialect_item: ($) =>
663
+ seq($.dialect_namespace, $.pretty_dialect_item_body),
664
+ pretty_dialect_item_body: ($) =>
665
+ seq("<", repeat($._pretty_dialect_item_contents), ">"),
666
+ _pretty_dialect_item_contents: ($) =>
667
+ prec.left(
668
+ choice(
669
+ $.pretty_dialect_item_body,
670
+ $._pretty_dialect_bang_body_token,
671
+ $._pretty_dialect_body_attribute,
672
+ $.dialect_dim_list,
673
+ $.type,
674
+ prec(2, $.attribute),
675
+ $._literal,
676
+ $._dense_keyword,
677
+ $._sparse_keyword,
678
+ "array",
679
+ "vector",
680
+ "tensor",
681
+ "opaque",
682
+ $.bare_id,
683
+ ",",
684
+ ":",
685
+ "=",
686
+ "->",
687
+ "(",
688
+ ")",
689
+ "[",
690
+ "]",
691
+ "{",
692
+ "}",
693
+ "*",
694
+ "?",
695
+ "@",
696
+ "#",
697
+ token(prec(-1, /[^<>]/)),
698
+ ),
699
+ ),
700
+ _pretty_dialect_bang_body_token: ($) =>
701
+ token(prec(1, seq("!", /[^a-zA-Z_<>]/))),
702
+ _pretty_dialect_body_attribute: ($) =>
703
+ choice(
704
+ alias($._pretty_dialect_attribute_token, $.dialect_attribute),
705
+ alias($._pretty_attribute_alias_token, $.attribute_alias),
706
+ ),
707
+ _pretty_attribute_alias_token: ($) =>
708
+ token(prec(2, seq("#", /[a-zA-Z_]/, repeat(/[a-zA-Z0-9_$]/)))),
709
+ _pretty_dialect_attribute_token: ($) =>
710
+ token(
711
+ prec(
712
+ 2,
713
+ seq(
714
+ "#",
715
+ /[a-zA-Z_]/,
716
+ repeat(/[a-zA-Z0-9_$]/),
717
+ ".",
718
+ /[a-zA-Z_]/,
719
+ repeat(/[a-zA-Z0-9_$.-]/),
720
+ ),
721
+ ),
722
+ ),
723
+ dimension_separator: ($) => token(prec(10, "x")),
724
+ dialect_dim_list: ($) =>
725
+ prec.dynamic(
726
+ 1,
727
+ seq(
728
+ $._dialect_dim_primitive,
729
+ repeat1(seq($.dimension_separator, $._dialect_dim_primitive)),
730
+ ),
731
+ ),
732
+ _dialect_dim_primitive: ($) =>
733
+ choice(
734
+ prec(1, $.type),
735
+ prec(1, $._pretty_dialect_body_type),
736
+ alias($.integer_literal, $.dimension_size),
737
+ "?",
738
+ "*",
739
+ ),
740
+ _pretty_dialect_body_type: ($) =>
741
+ seq(choice($.bare_id, "array"), $.pretty_dialect_item_body),
382
742
 
383
743
  // Builtin types
384
- builtin_type: $ => choice(
385
- $.integer_type,
386
- $.float_type,
387
- $.complex_type,
388
- $.index_type,
389
- $.memref_type,
390
- $.none_type,
391
- $.tensor_type,
392
- $.vector_type,
393
- $.tuple_type,
394
- $.opaque_type),
395
-
396
- integer_type: $ => token(prec(5, seq(choice('si', 'ui', 'i'), /[0-9]/, repeat(/[0-9]/)))),
397
- float_type: $ => token(prec(5, choice('f16', 'tf32', 'f32', 'f64', 'f80', 'f128', 'bf16',
398
- 'f4E2M1FN', 'f6E2M3FN', 'f6E3M2FN', 'f8E3M4', 'f8E4M3', 'f8E4M3FN', 'f8E4M3FNUZ',
399
- 'f8E4M3B11FNUZ', 'f8E5M2', 'f8E5M2FNUZ', 'f8E8M0FNU'))),
400
- index_type: $ => token(prec(5, 'index')),
401
- none_type: $ => token(prec(5, 'none')),
402
- complex_type: $ => seq(token('complex'), '<', $._prim_type, '>'),
403
- _prim_type: $ => choice($.integer_type, $.float_type, $.index_type,
404
- $.complex_type, $.none_type, $.memref_type, $.vector_type, $.tensor_type,
405
- $.tuple_type, $.opaque_type, $.dialect_type, $.type_alias),
406
-
407
- memref_type: $ => seq(token('memref'), '<',
408
- field('dimension_list', $.dim_list),
409
- optional(seq(',', $.attribute_value)),
410
- optional(seq(',', $.attribute_value)), '>'),
411
- dim_list: $ => seq($._dim_primitive, repeat(seq('x', $._dim_primitive))),
744
+ builtin_type: ($) =>
745
+ choice(
746
+ $.integer_type,
747
+ $.float_type,
748
+ $.complex_type,
749
+ $.index_type,
750
+ $.memref_type,
751
+ $.none_type,
752
+ $.tensor_type,
753
+ $.vector_type,
754
+ $.tuple_type,
755
+ $.opaque_type,
756
+ ),
757
+
758
+ integer_type: ($) =>
759
+ token(prec(5, seq(choice("si", "ui", "i"), /[0-9]/, repeat(/[0-9]/)))),
760
+ float_type: ($) =>
761
+ token(
762
+ prec(
763
+ 5,
764
+ choice(
765
+ "f16",
766
+ "tf32",
767
+ "f32",
768
+ "f64",
769
+ "f80",
770
+ "f128",
771
+ "bf16",
772
+ "f4E2M1FN",
773
+ "f6E2M3FN",
774
+ "f6E3M2FN",
775
+ "f8E3M4",
776
+ "f8E4M3",
777
+ "f8E4M3FN",
778
+ "f8E4M3FNUZ",
779
+ "f8E4M3B11FNUZ",
780
+ "f8E5M2",
781
+ "f8E5M2FNUZ",
782
+ "f8E8M0FNU",
783
+ ),
784
+ ),
785
+ ),
786
+ index_type: ($) => token(prec(5, "index")),
787
+ none_type: ($) => token(prec(5, "none")),
788
+ complex_type: ($) => seq($._complex_type_start, $._prim_type, ">"),
789
+ _prim_type: ($) =>
790
+ choice(
791
+ $.integer_type,
792
+ $.float_type,
793
+ $.index_type,
794
+ $.complex_type,
795
+ $.none_type,
796
+ $.memref_type,
797
+ $.vector_type,
798
+ $.tensor_type,
799
+ $.tuple_type,
800
+ $.opaque_type,
801
+ $.dialect_type,
802
+ $.type_alias,
803
+ ),
804
+
805
+ memref_type: ($) =>
806
+ seq(
807
+ token("memref"),
808
+ "<",
809
+ field("dimension_list", $.dim_list),
810
+ optional(seq(",", $.attribute_value)),
811
+ optional(seq(",", $.attribute_value)),
812
+ ">",
813
+ ),
814
+ dim_list: ($) => seq($._dim_primitive, repeat(seq("x", $._dim_primitive))),
412
815
  // '*' represents the unranked form for both memref and tensor types
413
816
  // (UnrankedMemRefType, UnrankedTensorType). Vector uses its own
414
817
  // vector_dim_list / _static_dim_list and disallows '*' by construction.
415
- dimension_size: $ => token(repeat1(/[0-9]/)),
416
- _dim_primitive: $ => choice(prec(1, $.type), $.dimension_size, '?', '*'),
818
+ dimension_size: ($) => token(repeat1(/[0-9]/)),
819
+ _dim_primitive: ($) => choice(prec(1, $.type), $.dimension_size, "?", "*"),
417
820
 
418
- tensor_type: $ => seq(token('tensor'), '<', $.dim_list,
419
- optional(seq(',', $.tensor_encoding)), '>'),
420
- tensor_encoding: $ => $.attribute_value,
821
+ tensor_type: ($) =>
822
+ seq(
823
+ token("tensor"),
824
+ "<",
825
+ $.dim_list,
826
+ optional(seq(",", $.tensor_encoding)),
827
+ ">",
828
+ ),
829
+ tensor_encoding: ($) => $.attribute_value,
421
830
 
422
- vector_type: $ => prec(1, seq(token('vector'), '<', repeat($.vector_dim_list), $._prim_type, '>')),
423
- vector_dim_list: $ => prec.left(choice(seq($._static_dim_list, 'x',
424
- optional(seq('[', $._static_dim_list, ']', 'x'))), seq('[', $._static_dim_list, ']', 'x'))),
425
- _static_dim_list: $ => seq($.dimension_size, repeat(seq('x', $.dimension_size))),
831
+ vector_type: ($) =>
832
+ prec(
833
+ 1,
834
+ seq(token("vector"), "<", repeat($.vector_dim_list), $._prim_type, ">"),
835
+ ),
836
+ vector_dim_list: ($) =>
837
+ prec.left(
838
+ choice(
839
+ seq(
840
+ $._static_dim_list,
841
+ "x",
842
+ optional(seq("[", $._static_dim_list, "]", "x")),
843
+ ),
844
+ seq("[", $._static_dim_list, "]", "x"),
845
+ ),
846
+ ),
847
+ _static_dim_list: ($) =>
848
+ seq($.dimension_size, repeat(seq("x", $.dimension_size))),
426
849
 
427
- tuple_type: $ => seq(token('tuple'), '<', optional(seq($.tuple_dim, repeat(seq(',', $.tuple_dim)))), '>'),
428
- tuple_dim: $ => $._prim_type,
850
+ tuple_type: ($) =>
851
+ seq(
852
+ token("tuple"),
853
+ "<",
854
+ optional(seq($.tuple_dim, repeat(seq(",", $.tuple_dim)))),
855
+ ">",
856
+ ),
857
+ tuple_dim: ($) => $._prim_type,
429
858
 
430
859
  // opaque-type ::= `opaque` `<` string-literal `,` string-literal `>`
431
860
  // e.g. opaque<"llvm", "struct<(i32, float)>">
432
- opaque_type: $ => seq(token('opaque'), '<', $.string_literal, ',', $.string_literal, '>'),
861
+ opaque_type: ($) =>
862
+ seq(token("opaque"), "<", $.string_literal, ",", $.string_literal, ">"),
433
863
 
434
864
  // =========================================================================
435
865
  // Attributes
436
866
  // attribute-entry ::= (bare-id | string-literal) `=` attribute-value
437
867
  // attribute-value ::= attribute-alias | dialect-attribute | builtin-attribute
438
868
  // =========================================================================
439
- attribute_entry: $ => choice(
440
- seq(choice($.bare_id, $.string_literal), optional(seq('=', $.attribute_value))),
441
- // Array-valued entry (e.g. {["op.name"]} in transform dialect)
442
- $._attribute_array,
443
- ),
444
- attribute_value: $ => choice($._attribute_array, $._attribute_value_nobracket),
445
- _attribute_array: $ => seq('[', optional($._attribute_array_element),
446
- repeat(seq(',', $._attribute_array_element)), ']'),
447
- _attribute_array_element: $ => choice($._attribute_array, $._attribute_value_nobracket),
448
- _attribute_value_nobracket: $ => choice($.attribute_alias, $.dialect_attribute,
449
- $.builtin_attribute, $.dictionary_attribute, $._literal_and_type, $.type,
450
- $.function_type, $._affine_map_like, $.symbol_ref_id, $.trailing_location,
451
- seq(choice($.attribute_alias, $.dialect_attribute, $.builtin_attribute), $._type_annotation)),
452
- attribute: $ => choice($.attribute_alias, $.dialect_attribute,
453
- $.builtin_attribute, $.dictionary_attribute),
869
+ attribute_entry: ($) =>
870
+ choice(
871
+ seq(
872
+ choice($.bare_id, $.string_literal),
873
+ optional(seq("=", $.attribute_value)),
874
+ ),
875
+ // Array-valued entry (e.g. {["op.name"]} in transform dialect)
876
+ $._attribute_array,
877
+ ),
878
+ attribute_value: ($) =>
879
+ choice($._attribute_array, $._attribute_value_nobracket),
880
+ _attribute_array: ($) =>
881
+ seq(
882
+ "[",
883
+ optional($._attribute_array_element),
884
+ repeat(seq(",", $._attribute_array_element)),
885
+ "]",
886
+ ),
887
+ _attribute_array_element: ($) =>
888
+ choice($._attribute_array, $._attribute_value_nobracket),
889
+ _attribute_value_nobracket: ($) =>
890
+ choice(
891
+ $.attribute_alias,
892
+ $.dialect_attribute,
893
+ $.builtin_attribute,
894
+ $.dictionary_attribute,
895
+ $._literal_and_type,
896
+ $.type,
897
+ $.function_type,
898
+ $._affine_map_like,
899
+ $.symbol_ref_id,
900
+ $.trailing_location,
901
+ seq(
902
+ choice($.attribute_alias, $.dialect_attribute, $.builtin_attribute),
903
+ $._type_annotation,
904
+ ),
905
+ ),
906
+ attribute: ($) =>
907
+ choice(
908
+ $.attribute_alias,
909
+ $.dialect_attribute,
910
+ $.builtin_attribute,
911
+ $.dictionary_attribute,
912
+ ),
454
913
 
455
914
  // Attribute aliases
456
- attribute_alias_def: $ => seq('#', $._alias_or_dialect_id, '=', $.attribute_value),
457
- attribute_alias: $ => seq('#', $._alias_or_dialect_id),
915
+ attribute_alias_def: ($) =>
916
+ seq("#", $._alias_or_dialect_id, "=", $.attribute_value),
917
+ attribute_alias: ($) => seq("#", $._alias_or_dialect_id),
458
918
 
459
919
  // Dialect attributes
460
- dialect_attribute: $ => seq('#', choice($.opaque_dialect_item, $.pretty_dialect_item, $.parametric_dialect_item)),
920
+ dialect_attribute: ($) =>
921
+ seq(
922
+ "#",
923
+ choice(
924
+ $.opaque_dialect_item,
925
+ $.pretty_dialect_item,
926
+ $.parametric_dialect_item,
927
+ ),
928
+ ),
461
929
 
462
930
  // Builtin attributes
463
- builtin_attribute: $ => choice(
464
- $.strided_layout,
465
- prec(1, $.affine_map),
466
- prec(1, $.affine_set),
467
- $.dense_resource_literal,
468
- $.distinct_attribute,
469
- ),
470
- dense_resource_literal: $ => seq(token('dense_resource'), '<',
471
- choice($.bare_id, $.string_literal), '>'),
472
- distinct_attribute: $ => seq(token('distinct'), '[',
473
- alias($._unsigned_integer_literal, $.integer_literal), ']',
474
- '<', optional($.attribute_value), '>'),
475
- strided_layout: $ => seq(token('strided'), '<', '[', optional($._stride_list_comma), ']',
476
- optional(seq(',', token('offset'), ':', $._stride_primitive)), '>'),
477
- _stride_list_comma: $ => seq($._stride_primitive, repeat(seq(',', $._stride_primitive))),
478
- _stride_primitive: $ => choice($.integer_literal, '?', '*'),
931
+ builtin_attribute: ($) =>
932
+ choice(
933
+ $.strided_layout,
934
+ prec(1, $.affine_map),
935
+ prec(1, $.affine_set),
936
+ $.dense_resource_literal,
937
+ $.distinct_attribute,
938
+ ),
939
+ dense_resource_literal: ($) =>
940
+ seq(
941
+ token(prec(2, "dense_resource")),
942
+ "<",
943
+ choice($.bare_id, $.string_literal),
944
+ ">",
945
+ ),
946
+ distinct_attribute: ($) =>
947
+ seq(
948
+ token("distinct"),
949
+ "[",
950
+ alias($._unsigned_integer_literal, $.integer_literal),
951
+ "]",
952
+ "<",
953
+ optional($.attribute_value),
954
+ ">",
955
+ ),
956
+ strided_layout: ($) =>
957
+ seq(
958
+ token("strided"),
959
+ "<",
960
+ "[",
961
+ optional($._stride_list_comma),
962
+ "]",
963
+ optional(seq(",", token("offset"), ":", $._stride_primitive)),
964
+ ">",
965
+ ),
966
+ _stride_list_comma: ($) =>
967
+ seq($._stride_primitive, repeat(seq(",", $._stride_primitive))),
968
+ _stride_primitive: ($) => choice($.integer_literal, "?", "*"),
479
969
 
480
970
  // =========================================================================
481
971
  // Affine expressions
482
972
  // =========================================================================
483
- affine_map: $ => seq(token('affine_map'), '<', $._multi_dim_affine_expr_parens,
484
- optional($._multi_dim_affine_expr_sq), token('->'), $._multi_dim_affine_expr_parens, '>'),
485
- _affine_map_like: $ => seq($._multi_dim_affine_expr_parens,
486
- optional($._multi_dim_affine_expr_sq), token('->'), $._multi_dim_affine_expr_parens),
487
- affine_set: $ => seq(token('affine_set'), '<', $._multi_dim_affine_expr_parens,
488
- optional($._multi_dim_affine_expr_sq), ':', $._multi_dim_affine_expr_parens, '>'),
489
- _multi_dim_affine_expr_parens: $ => seq('(', optional($._multi_dim_affine_expr), ')'),
490
- _multi_dim_affine_expr_sq: $ => seq('[', optional($._multi_dim_affine_expr), ']'),
491
-
492
- _multi_dim_affine_expr: $ => seq($._affine_expr, repeat(seq(',', $._affine_expr))),
493
- _affine_expr: $ => choice(
494
- seq('(', $._affine_expr, ')'),
495
- prec(4, seq('-', $._affine_expr)),
496
- prec.left(3, seq($._affine_expr, choice('*', 'ceildiv', 'floordiv', 'mod'), $._affine_expr)),
497
- prec.left(2, seq($._affine_expr, choice('+', '-'), $._affine_expr)),
498
- prec.left(1, seq($._affine_expr, choice('==', '>=', '<='), $._affine_expr)),
499
- $._affine_prim
500
- ),
501
- _affine_prim: $ => choice($.integer_literal, $.value_use,
502
- seq(choice($.bare_id, 'dense', 'sparse'), optional(seq(':', choice($.bare_id, 'dense', 'sparse', 'compressed', 'singleton', 'loose_compressed', 'n_out_of_m')))),
503
- seq('symbol', '(', $.value_use, ')'), seq(choice('max', 'min'), '(', $._value_use_list, ')')),
973
+ affine_map: ($) =>
974
+ seq(
975
+ token("affine_map"),
976
+ "<",
977
+ $._multi_dim_affine_expr_parens,
978
+ optional($._multi_dim_affine_expr_sq),
979
+ token("->"),
980
+ $._multi_dim_affine_expr_parens,
981
+ ">",
982
+ ),
983
+ _affine_map_like: ($) =>
984
+ seq(
985
+ $._multi_dim_affine_expr_parens,
986
+ optional($._multi_dim_affine_expr_sq),
987
+ token("->"),
988
+ $._multi_dim_affine_expr_parens,
989
+ ),
990
+ affine_set: ($) =>
991
+ seq(
992
+ token("affine_set"),
993
+ "<",
994
+ $._multi_dim_affine_expr_parens,
995
+ optional($._multi_dim_affine_expr_sq),
996
+ ":",
997
+ $._multi_dim_affine_expr_parens,
998
+ ">",
999
+ ),
1000
+ _multi_dim_affine_expr_parens: ($) =>
1001
+ seq("(", optional($._multi_dim_affine_expr), ")"),
1002
+ _multi_dim_affine_expr_sq: ($) =>
1003
+ seq("[", optional($._multi_dim_affine_expr), "]"),
1004
+
1005
+ _multi_dim_affine_expr: ($) =>
1006
+ seq($._affine_expr, repeat(seq(",", $._affine_expr))),
1007
+ _affine_expr: ($) =>
1008
+ choice(
1009
+ seq("(", $._affine_expr, ")"),
1010
+ prec(4, seq("-", $._affine_expr)),
1011
+ prec.left(
1012
+ 3,
1013
+ seq(
1014
+ $._affine_expr,
1015
+ choice("*", "ceildiv", "floordiv", "mod"),
1016
+ $._affine_expr,
1017
+ ),
1018
+ ),
1019
+ prec.left(2, seq($._affine_expr, choice("+", "-"), $._affine_expr)),
1020
+ prec.left(
1021
+ 1,
1022
+ seq($._affine_expr, choice("==", ">=", "<="), $._affine_expr),
1023
+ ),
1024
+ $._affine_prim,
1025
+ ),
1026
+ _affine_prim: ($) =>
1027
+ choice(
1028
+ $.integer_literal,
1029
+ $.value_use,
1030
+ seq(
1031
+ choice($.bare_id, $._dense_keyword, $._sparse_keyword),
1032
+ optional(
1033
+ seq(
1034
+ ":",
1035
+ choice(
1036
+ $.bare_id,
1037
+ $._dense_keyword,
1038
+ $._sparse_keyword,
1039
+ "compressed",
1040
+ "singleton",
1041
+ "loose_compressed",
1042
+ "n_out_of_m",
1043
+ ),
1044
+ ),
1045
+ ),
1046
+ ),
1047
+ seq("symbol", "(", $.value_use, ")"),
1048
+ seq(choice("max", "min"), "(", $._value_use_list, ")"),
1049
+ ),
504
1050
 
505
1051
  // =========================================================================
506
1052
  // Function-related rules (used by func_operation tier-1)
507
1053
  // =========================================================================
508
- func_return: $ => seq(token('->'), $.type_list_attr_parens),
509
- func_arg_list: $ => seq('(', optional(choice($.variadic,
510
- $._value_id_and_type_attr_list)), ')'),
511
- _value_id_and_type_attr_list: $ => seq($._value_id_and_type_attr,
512
- repeat(seq(',', choice($._value_id_and_type_attr, $.variadic)))),
513
- _value_id_and_type_attr: $ => seq($._function_arg, optional($.attribute),
514
- optional($.trailing_location)),
515
- _function_arg: $ => choice(seq($.value_use, ':', choice($.type, $.function_type)), $.value_use, $.type, $.function_type),
516
- _type_or_func_type: $ => choice($.type, $.function_type),
517
- type_list_attr_parens: $ => choice($.type, $.function_type,
518
- seq('(', $._type_or_func_type, optional($.attribute),
519
- repeat(seq(',', $._type_or_func_type, optional($.attribute))), ')'), seq('(', ')')),
520
- variadic: $ => token('...'),
1054
+ func_return: ($) => seq(token("->"), $.type_list_attr_parens),
1055
+ func_arg_list: ($) =>
1056
+ seq(
1057
+ "(",
1058
+ optional(choice($.variadic, $._value_id_and_type_attr_list)),
1059
+ ")",
1060
+ ),
1061
+ _value_id_and_type_attr_list: ($) =>
1062
+ seq(
1063
+ $._value_id_and_type_attr,
1064
+ repeat(seq(",", choice($._value_id_and_type_attr, $.variadic))),
1065
+ ),
1066
+ _value_id_and_type_attr: ($) =>
1067
+ seq(
1068
+ $._function_arg,
1069
+ optional($.attribute),
1070
+ optional($.trailing_location),
1071
+ ),
1072
+ _function_arg: ($) =>
1073
+ choice(
1074
+ seq($.value_use, ":", choice($.type, $.function_type)),
1075
+ $.value_use,
1076
+ $.type,
1077
+ $.function_type,
1078
+ ),
1079
+ _type_or_func_type: ($) => choice($.type, $.function_type),
1080
+ type_list_attr_parens: ($) =>
1081
+ choice(
1082
+ $.type,
1083
+ $.function_type,
1084
+ seq(
1085
+ "(",
1086
+ $._type_or_func_type,
1087
+ optional($.attribute),
1088
+ repeat(seq(",", $._type_or_func_type, optional($.attribute))),
1089
+ ")",
1090
+ ),
1091
+ seq("(", ")"),
1092
+ ),
1093
+ variadic: ($) => token("..."),
521
1094
 
522
1095
  // =========================================================================
523
1096
  // Shared helpers
524
1097
  // =========================================================================
525
- _value_use_list_parens: $ => seq('(', optional($._value_use_list), ')'),
1098
+ _value_use_list_parens: ($) => seq("(", optional($._value_use_list), ")"),
526
1099
 
527
1100
  // Comment (standard BCPL)
528
- comment: $ => token(seq('//', /.*/)),
529
- }
1101
+ comment: ($) => token(seq("//", /.*/)),
1102
+ _dense_keyword: ($) => token(prec(1, "dense")),
1103
+ _sparse_keyword: ($) => token(prec(1, "sparse")),
1104
+ _complex_label_start: ($) => token(prec(2, /complex:/)),
1105
+ _complex_type_start: ($) => token(prec(1, /complex</)),
1106
+ },
530
1107
  });