1 /*
2 * Copyright (c) 2025, 2026, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25
26 package oracle.code.onnx.ir;
27
28 import java.util.*;
29 import jdk.incubator.code.*;
30 import jdk.incubator.code.Op.Nested;
31 import jdk.incubator.code.extern.ExternalizedOp;
32
33 public sealed class ExplicitOnnxOps permits OnnxOps {
34
35 // @@@ this should be generated from contrib operators
36 @OpFactoryHelper.OpDeclaration(GroupQueryAttention.NAME)
37 public static final class GroupQueryAttention extends OnnxOp {
38 public static final String NAME = "com.microsoft.GroupQueryAttention";
39
40 public enum Attribute implements OnnxAttribute {
41 do_rotary(Long.class, true, 0),
42 kv_num_heads(Long.class, false, null),
43 local_window_size(Long.class, true, -1),
44 num_heads(Long.class, false, null),
45 rotary_interleaved(Long.class, true, 0),
46 scale(Float.class, true, null), // @@@ Default value is 1/sqrt(head_size)
47 ;
48
49 final Class<?> t;
50 final boolean optional;
51 final Object defaultValue;
52
53 Attribute(Class<?> type, boolean optional, Object defaultValue) {
54 this.t = type;
55 this.optional = optional;
56 this.defaultValue = defaultValue;
57 assert optional || defaultValue == null;
58 }
59
60 public Class<?> type() {
61 return t;
62 }
63
64 public boolean isOptional() {
65 return optional;
66 }
67
68 public Object defaultValue() {
69 return defaultValue;
70 }
71 }
72
73 public enum TypeConstraint implements OnnxTypeConstraint {
74 T(new OnnxType.TypeVariable("T", List.of(OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.bfloat16()), OnnxType.tensor(OnnxType.float32())))),
75 M(new OnnxType.TypeVariable("M", List.of(OnnxType.tensor(OnnxType.int32())))),
76 ;
77
78 final OnnxType.TypeVariable typeVariable;
79
80 TypeConstraint(OnnxType.TypeVariable typeVariable) {
81 assert typeVariable.name().equals(name());
82 this.typeVariable = typeVariable;
83 }
84
85 @Override
86 public OnnxType.TypeVariable typeVariable() {
87 return typeVariable;
88 }
89 }
90
91 public enum InputParameter implements OnnxParameter {
92 query(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
93 key(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
94 value(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
95 past_key(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
96 past_value(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
97 seqlens_k(TypeConstraint.M.typeVariable(), Quantifier.REQUIRED),
98 total_sequence_length(TypeConstraint.M.typeVariable(), Quantifier.REQUIRED),
99 cos_cache(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
100 sin_cache(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
101 ;
102
103 final OnnxType type;
104 final Quantifier quantifier;
105
106 InputParameter(OnnxType type, Quantifier quantifier) {
107 this.type = type;
108 this.quantifier = quantifier;
109 }
110
111 @Override
112 public OnnxType type() {
113 return type;
114 }
115
116 @Override
117 public Quantifier quantifier() {
118 return quantifier;
119 }
120 }
121
122 public enum OutputParameter implements OnnxParameter {
123 output(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
124 present_key(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
125 present_value(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
126 ;
127
128 final OnnxType type;
129 final Quantifier quantifier;
130
131 OutputParameter(OnnxType type, Quantifier quantifier) {
132 this.type = type;
133 this.quantifier = quantifier;
134 }
135
136 @Override
137 public OnnxType type() {
138 return type;
139 }
140
141 @Override
142 public Quantifier quantifier() {
143 return quantifier;
144 }
145 }
146
147 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
148 NAME,
149 List.of(Attribute.values()),
150 List.of(TypeConstraint.values()),
151 List.of(InputParameter.values()),
152 List.of(OutputParameter.values())
153 );
154
155 public GroupQueryAttention(ExternalizedOp def) {
156 super(SCHEMA, def);
157 }
158
159 GroupQueryAttention(GroupQueryAttention that, CodeContext cc) {
160 super(that, cc);
161 }
162
163 @Override
164 public GroupQueryAttention transform(CodeContext cc, CodeTransformer ot) {
165 return new GroupQueryAttention(this, cc);
166 }
167
168 GroupQueryAttention(CodeType resultType, Value query, java.util.Optional<Value> key, java.util.Optional<Value> value, java.util.Optional<Value> past_key, java.util.Optional<Value> past_value, Value seqlens_k, Value total_sequence_length, java.util.Optional<Value> cos_cache, java.util.Optional<Value> sin_cache, java.util.Optional<Long> do_rotary, long kv_num_heads, java.util.Optional<Long> local_window_size, long num_heads, java.util.Optional<Long> rotary_interleaved, java.util.Optional<Float> scale) {
169 super(SCHEMA, resultType, Collections.emptySet(), List.of(query, key, value, past_key, past_value, seqlens_k, total_sequence_length, cos_cache, sin_cache), List.of(do_rotary, kv_num_heads, local_window_size, num_heads, rotary_interleaved, scale));
170 }
171
172 @Override
173 public SequencedSet<OnnxParameter> onnxOutputs() {
174 return onnxOutputs(SCHEMA);
175 }
176
177 @Override
178 public SequencedMap<OnnxParameter, Object> onnxInputs() {
179 return onnxInputs(SCHEMA, List.of(query(), key(), value(), past_key(), past_value(), seqlens_k(), total_sequence_length(), cos_cache(), sin_cache()));
180 }
181
182 public Value query() {
183 return operands().get(0);
184 }
185
186 public java.util.Optional<Value> key() {
187 int i = optionalInputArguments.indexOf(InputParameter.key);
188 return i != -1 ? java.util.Optional.of(operands().get(1 + i)) : java.util.Optional.empty();
189 }
190
191 public java.util.Optional<Value> value() {
192 int i = optionalInputArguments.indexOf(InputParameter.value);
193 return i != -1 ? java.util.Optional.of(operands().get(1 + i)) : java.util.Optional.empty();
194 }
195
196 public java.util.Optional<Value> past_key() {
197 int i = optionalInputArguments.indexOf(InputParameter.past_key);
198 return i != -1 ? java.util.Optional.of(operands().get(1 + i)) : java.util.Optional.empty();
199 }
200
201 public java.util.Optional<Value> past_value() {
202 int i = optionalInputArguments.indexOf(InputParameter.past_value);
203 return i != -1 ? java.util.Optional.of(operands().get(1 + i)) : java.util.Optional.empty();
204 }
205
206 private int skipOptional() {
207 for (int i = optionalInputArguments.size() - 1; i >= 0; i--) {
208 var opt = optionalInputArguments.get(i);
209 if (opt != InputParameter.cos_cache && opt != InputParameter.sin_cache) return i;
210 }
211 return -1;
212 }
213
214 public Value seqlens_k() {
215 return operands().get(skipOptional() + 2);
216 }
217
218 public Value total_sequence_length() {
219 return operands().get(skipOptional() + 3);
220 }
221
222 public java.util.Optional<Value> cos_cache() {
223 int i = optionalInputArguments.indexOf(InputParameter.cos_cache);
224 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
225 }
226
227 public java.util.Optional<Value> sin_cache() {
228 int i = optionalInputArguments.indexOf(InputParameter.sin_cache);
229 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
230 }
231 }
232
233 public static GroupQueryAttention GroupQueryAttention(CodeType resultType, Value query, java.util.Optional<Value> key, java.util.Optional<Value> value, java.util.Optional<Value> past_key, java.util.Optional<Value> past_value, Value seqlens_k, Value total_sequence_length, java.util.Optional<Value> cos_cache, java.util.Optional<Value> sin_cache, java.util.Optional<Long> do_rotary, long kv_num_heads, java.util.Optional<Long> local_window_size, long num_heads, java.util.Optional<Long> rotary_interleaved, java.util.Optional<Float> scale) {
234 return new GroupQueryAttention(resultType, query, key, value, past_key, past_value, seqlens_k, total_sequence_length, cos_cache, sin_cache, do_rotary, kv_num_heads, local_window_size, num_heads, rotary_interleaved, scale);
235 }
236
237 // @@@ this should be generated from contrib operators
238 @OpFactoryHelper.OpDeclaration(MatMulNBits.NAME)
239 public static final class MatMulNBits extends OnnxOp {
240 public static final String NAME = "com.microsoft.MatMulNBits";
241
242 public enum Attribute implements OnnxAttribute {
243 K(Long.class, false, null),
244 N(Long.class, false, null),
245 accuracy_level(Long.class, true, 0),
246 bits(Long.class, false, null),
247 block_size(Long.class, false, null),
248 ;
249
250 final Class<?> t;
251 final boolean optional;
252 final Object defaultValue;
253
254 Attribute(Class<?> type, boolean optional, Object defaultValue) {
255 this.t = type;
256 this.optional = optional;
257 this.defaultValue = defaultValue;
258 assert optional || defaultValue == null;
259 }
260
261 public Class<?> type() {
262 return t;
263 }
264
265 public boolean isOptional() {
266 return optional;
267 }
268
269 public Object defaultValue() {
270 return defaultValue;
271 }
272 }
273
274 public enum TypeConstraint implements OnnxTypeConstraint {
275 T1(new OnnxType.TypeVariable("T1", List.of(OnnxType.tensor(OnnxType.float32()), OnnxType.tensor(OnnxType.float16())))),
276 T2(new OnnxType.TypeVariable("T2", List.of(OnnxType.tensor(OnnxType.uint8()), OnnxType.tensor(OnnxType.int32())))),
277 T3(new OnnxType.TypeVariable("T3", List.of(OnnxType.tensor(OnnxType.uint8()), OnnxType.tensor(OnnxType.int32()), OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.float32())))),
278 T4(new OnnxType.TypeVariable("T4", List.of(OnnxType.tensor(OnnxType.int32())))),
279 ;
280
281 final OnnxType.TypeVariable typeVariable;
282
283 TypeConstraint(OnnxType.TypeVariable typeVariable) {
284 assert typeVariable.name().equals(name());
285 this.typeVariable = typeVariable;
286 }
287
288 @Override
289 public OnnxType.TypeVariable typeVariable() {
290 return typeVariable;
291 }
292 }
293
294 public enum InputParameter implements OnnxParameter {
295 A(TypeConstraint.T1.typeVariable(), Quantifier.REQUIRED),
296 B(TypeConstraint.T2.typeVariable(), Quantifier.REQUIRED),
297 scales(TypeConstraint.T1.typeVariable(), Quantifier.REQUIRED),
298 zero_points(TypeConstraint.T3.typeVariable(), Quantifier.OPTIONAL),
299 g_idx(TypeConstraint.T4.typeVariable(), Quantifier.OPTIONAL),
300 bias(TypeConstraint.T1.typeVariable(), Quantifier.OPTIONAL),
301 ;
302
303 final OnnxType type;
304 final Quantifier quantifier;
305
306 InputParameter(OnnxType type, Quantifier quantifier) {
307 this.type = type;
308 this.quantifier = quantifier;
309 }
310
311 @Override
312 public OnnxType type() {
313 return type;
314 }
315
316 @Override
317 public Quantifier quantifier() {
318 return quantifier;
319 }
320 }
321
322 public enum OutputParameter implements OnnxParameter {
323 Y(TypeConstraint.T1.typeVariable(), Quantifier.REQUIRED),
324 ;
325
326 final OnnxType type;
327 final Quantifier quantifier;
328
329 OutputParameter(OnnxType type, Quantifier quantifier) {
330 this.type = type;
331 this.quantifier = quantifier;
332 }
333
334 @Override
335 public OnnxType type() {
336 return type;
337 }
338
339 @Override
340 public Quantifier quantifier() {
341 return quantifier;
342 }
343 }
344
345 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
346 NAME,
347 List.of(Attribute.values()),
348 List.of(TypeConstraint.values()),
349 List.of(InputParameter.values()),
350 List.of(OutputParameter.values())
351 );
352
353 public MatMulNBits(ExternalizedOp def) {
354 super(SCHEMA, def);
355 }
356
357 MatMulNBits(MatMulNBits that, CodeContext cc) {
358 super(that, cc);
359 }
360
361 @Override
362 public MatMulNBits transform(CodeContext cc, CodeTransformer ot) {
363 return new MatMulNBits(this, cc);
364 }
365
366 MatMulNBits(CodeType resultType, Value a, Value b, Value scales, java.util.Optional<Value> zero_points, java.util.Optional<Value> g_idx, java.util.Optional<Value> bias, long K, long N, java.util.Optional<Long> accuracy_level, long bits, long block_size) {
367 super(SCHEMA, resultType, Collections.emptySet(), List.of(a, b, scales, zero_points, g_idx, bias), List.of(K, N, accuracy_level, bits, block_size));
368 }
369
370 @Override
371 public SequencedSet<OnnxParameter> onnxOutputs() {
372 return onnxOutputs(SCHEMA);
373 }
374
375 @Override
376 public SequencedMap<OnnxParameter, Object> onnxInputs() {
377 return onnxInputs(SCHEMA, List.of(a(), b(), scales(), zero_points(), g_idx(), bias()));
378 }
379
380 public Value a() {
381 return operands().get(0);
382 }
383
384 public Value b() {
385 return operands().get(1);
386 }
387
388 public Value scales() {
389 return operands().get(2);
390 }
391
392 public java.util.Optional<Value> zero_points() {
393 int i = optionalInputArguments.indexOf(InputParameter.zero_points);
394 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
395 }
396
397 public java.util.Optional<Value> g_idx() {
398 int i = optionalInputArguments.indexOf(InputParameter.g_idx);
399 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
400 }
401
402 public java.util.Optional<Value> bias() {
403 int i = optionalInputArguments.indexOf(InputParameter.bias);
404 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
405 }
406 }
407
408 // @@@ this should be generated from contrib operators
409 @OpFactoryHelper.OpDeclaration(MultiHeadAttention.NAME)
410 public static final class MultiHeadAttention extends OnnxOp {
411 public static final String NAME = "com.microsoft.MultiHeadAttention";
412
413 public enum Attribute implements OnnxAttribute {
414 mask_filter_value(Float.class, true, null),
415 num_heads(Long.class, false, null),
416 scale(Float.class, true, null),
417 unidirectional(Long.class, true, 0),
418 ;
419
420 final Class<?> t;
421 final boolean optional;
422 final Object defaultValue;
423
424 Attribute(Class<?> type, boolean optional, Object defaultValue) {
425 this.t = type;
426 this.optional = optional;
427 this.defaultValue = defaultValue;
428 assert optional || defaultValue == null;
429 }
430
431 public Class<?> type() {
432 return t;
433 }
434
435 public boolean isOptional() {
436 return optional;
437 }
438
439 public Object defaultValue() {
440 return defaultValue;
441 }
442 }
443
444 public enum TypeConstraint implements OnnxTypeConstraint {
445 T(new OnnxType.TypeVariable("T", List.of(OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.float32())))),
446 M(new OnnxType.TypeVariable("M", List.of(OnnxType.tensor(OnnxType.int32())))),
447 ;
448
449 final OnnxType.TypeVariable typeVariable;
450
451 TypeConstraint(OnnxType.TypeVariable typeVariable) {
452 assert typeVariable.name().equals(name());
453 this.typeVariable = typeVariable;
454 }
455
456 @Override
457 public OnnxType.TypeVariable typeVariable() {
458 return typeVariable;
459 }
460 }
461
462 public enum InputParameter implements OnnxParameter {
463 query(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
464 key(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
465 value(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
466 bias(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
467 key_padding_mask(TypeConstraint.M.typeVariable(), Quantifier.OPTIONAL),
468 relative_position_bias(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
469 past_key(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
470 past_value(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
471 attention_bias(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
472 ;
473
474 final OnnxType type;
475 final Quantifier quantifier;
476
477 InputParameter(OnnxType type, Quantifier quantifier) {
478 this.type = type;
479 this.quantifier = quantifier;
480 }
481
482 @Override
483 public OnnxType type() {
484 return type;
485 }
486
487 @Override
488 public Quantifier quantifier() {
489 return quantifier;
490 }
491 }
492
493 public enum OutputParameter implements OnnxParameter {
494 output(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
495 present_key(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
496 present_value(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
497 ;
498
499 final OnnxType type;
500 final Quantifier quantifier;
501
502 OutputParameter(OnnxType type, Quantifier quantifier) {
503 this.type = type;
504 this.quantifier = quantifier;
505 }
506
507 @Override
508 public OnnxType type() {
509 return type;
510 }
511
512 @Override
513 public Quantifier quantifier() {
514 return quantifier;
515 }
516 }
517
518 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
519 NAME,
520 List.of(Attribute.values()),
521 List.of(TypeConstraint.values()),
522 List.of(InputParameter.values()),
523 List.of(OutputParameter.values())
524 );
525
526 public MultiHeadAttention(ExternalizedOp def) {
527 super(SCHEMA, def);
528 }
529
530 MultiHeadAttention(MultiHeadAttention that, CodeContext cc) {
531 super(that, cc);
532 }
533
534 @Override
535 public MultiHeadAttention transform(CodeContext cc, CodeTransformer ot) {
536 return new MultiHeadAttention(this, cc);
537 }
538
539 MultiHeadAttention(CodeType resultType, Set<OutputParameter> optionalOutputs, Value query, Value key, Value value, java.util.Optional<Value> bias, java.util.Optional<Value> key_padding_mask, java.util.Optional<Value> relative_position_bias, java.util.Optional<Value> past_key, java.util.Optional<Value> past_value, java.util.Optional<Value> attention_bias, long num_heads, java.util.Optional<Float> mask_filter_value, java.util.Optional<Float> scale, java.util.Optional<Long> unidirectional) {
540 super(SCHEMA, resultType, optionalOutputs, List.of(query, key, value, bias, key_padding_mask, relative_position_bias, past_key, past_value, attention_bias), List.of(mask_filter_value, num_heads, scale, unidirectional));
541 }
542
543 @Override
544 public SequencedSet<OnnxParameter> onnxOutputs() {
545 return onnxOutputs(SCHEMA);
546 }
547
548 @Override
549 public SequencedMap<OnnxParameter, Object> onnxInputs() {
550 return onnxInputs(SCHEMA, List.of(query(), key(), value(), bias(), key_padding_mask(), relative_position_bias(), past_key(), past_value(), attention_bias()));
551 }
552
553 public Value query() {
554 return operands().get(0);
555 }
556
557 public Value key() {
558 return operands().get(1);
559 }
560
561 public Value value() {
562 return operands().get(2);
563 }
564
565 public java.util.Optional<Value> bias() {
566 int i = optionalInputArguments.indexOf(InputParameter.bias);
567 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
568 }
569
570 public java.util.Optional<Value> key_padding_mask() {
571 int i = optionalInputArguments.indexOf(InputParameter.key_padding_mask);
572 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
573 }
574
575 public java.util.Optional<Value> relative_position_bias() {
576 int i = optionalInputArguments.indexOf(InputParameter.relative_position_bias);
577 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
578 }
579
580 public java.util.Optional<Value> past_key() {
581 int i = optionalInputArguments.indexOf(InputParameter.past_key);
582 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
583 }
584
585 public java.util.Optional<Value> past_value() {
586 int i = optionalInputArguments.indexOf(InputParameter.past_value);
587 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
588 }
589
590 public java.util.Optional<Value> attention_bias() {
591 int i = optionalInputArguments.indexOf(InputParameter.attention_bias);
592 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
593 }
594 }
595
596 public static MultiHeadAttention MultiHeadAttention(CodeType resultType, Set<MultiHeadAttention.OutputParameter> optionalOutputs, Value query, Value key, Value value, java.util.Optional<Value> bias, java.util.Optional<Value> key_padding_mask, java.util.Optional<Value> relative_position_bias, java.util.Optional<Value> past_key, java.util.Optional<Value> past_value, java.util.Optional<Value> attention_bias, long num_heads, java.util.Optional<Float> mask_filter_value, java.util.Optional<Float> scale, java.util.Optional<Long> unidirectional) {
597 return new MultiHeadAttention(resultType, optionalOutputs, query, key, value, bias, key_padding_mask, relative_position_bias, past_key, past_value, attention_bias, num_heads, mask_filter_value, scale, unidirectional);
598 }
599
600 // @@@ this should be generated from contrib operators
601 @OpFactoryHelper.OpDeclaration(FastGelu.NAME)
602 public static final class FastGelu extends OnnxOp {
603 public static final String NAME = "com.microsoft.FastGelu";
604
605 public enum TypeConstraint implements OnnxTypeConstraint {
606 T(new OnnxType.TypeVariable("T", List.of(OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.float32())))),
607 ;
608
609 final OnnxType.TypeVariable typeVariable;
610
611 TypeConstraint(OnnxType.TypeVariable typeVariable) {
612 assert typeVariable.name().equals(name());
613 this.typeVariable = typeVariable;
614 }
615
616 @Override
617 public OnnxType.TypeVariable typeVariable() {
618 return typeVariable;
619 }
620 }
621
622 public enum InputParameter implements OnnxParameter {
623 X(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
624 bias(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
625 ;
626
627 final OnnxType type;
628 final Quantifier quantifier;
629
630 InputParameter(OnnxType type, Quantifier quantifier) {
631 this.type = type;
632 this.quantifier = quantifier;
633 }
634
635 @Override
636 public OnnxType type() {
637 return type;
638 }
639
640 @Override
641 public Quantifier quantifier() {
642 return quantifier;
643 }
644 }
645
646 public enum OutputParameter implements OnnxParameter {
647 Y(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
648 ;
649
650 final OnnxType type;
651 final Quantifier quantifier;
652
653 OutputParameter(OnnxType type, Quantifier quantifier) {
654 this.type = type;
655 this.quantifier = quantifier;
656 }
657
658 @Override
659 public OnnxType type() {
660 return type;
661 }
662
663 @Override
664 public Quantifier quantifier() {
665 return quantifier;
666 }
667 }
668
669 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
670 NAME,
671 List.of(),
672 List.of(TypeConstraint.values()),
673 List.of(InputParameter.values()),
674 List.of(OutputParameter.values())
675 );
676
677 public FastGelu(ExternalizedOp def) {
678 super(SCHEMA, def);
679 }
680
681 FastGelu(FastGelu that, CodeContext cc) {
682 super(that, cc);
683 }
684
685 @Override
686 public FastGelu transform(CodeContext cc, CodeTransformer ot) {
687 return new FastGelu(this, cc);
688 }
689
690 FastGelu(CodeType resultType, Value X, java.util.Optional<Value> bias) {
691 super(SCHEMA, resultType, Collections.emptySet(), List.of(X, bias), List.of());
692 }
693
694 @Override
695 public SequencedSet<OnnxParameter> onnxOutputs() {
696 return onnxOutputs(SCHEMA);
697 }
698
699 @Override
700 public SequencedMap<OnnxParameter, Object> onnxInputs() {
701 return onnxInputs(SCHEMA, List.of(X(), bias()));
702 }
703
704 public Value X() {
705 return operands().get(0);
706 }
707
708 public java.util.Optional<Value> bias() {
709 int i = optionalInputArguments.indexOf(InputParameter.bias);
710 return i != -1 ? java.util.Optional.of(operands().get(1 + i)) : java.util.Optional.empty();
711 }
712 }
713
714 public static FastGelu FastGelu(CodeType resultType, Value X, java.util.Optional<Value> bias) {
715 return new FastGelu(resultType, X, bias);
716 }
717
718 public static MatMulNBits MatMulNBits(CodeType resultType, Value a, Value b, Value scales, java.util.Optional<Value> zero_points, java.util.Optional<Value> g_idx, java.util.Optional<Value> bias, long K, long N, java.util.Optional<Long> accuracy_level, long bits, long block_size) {
719 return new MatMulNBits(resultType, a, b, scales, zero_points, g_idx, bias, K, N, accuracy_level, bits, block_size);
720 }
721
722 // @@@ this should be generated from contrib operators
723 @OpFactoryHelper.OpDeclaration(SkipSimplifiedLayerNormalization.NAME)
724 public static final class SkipSimplifiedLayerNormalization extends OnnxOp {
725 public static final String NAME = "com.microsoft.SkipSimplifiedLayerNormalization";
726
727 public enum Attribute implements OnnxAttribute {
728 epsilon(Float.class, true, null),
729 ;
730
731 final Class<?> t;
732 final boolean optional;
733 final Object defaultValue;
734
735 Attribute(Class<?> type, boolean optional, Object defaultValue) {
736 this.t = type;
737 this.optional = optional;
738 this.defaultValue = defaultValue;
739 assert optional || defaultValue == null;
740 }
741
742 public Class<?> type() {
743 return t;
744 }
745
746 public boolean isOptional() {
747 return optional;
748 }
749
750 public Object defaultValue() {
751 return defaultValue;
752 }
753 }
754
755 public enum TypeConstraint implements OnnxTypeConstraint {
756 T(new OnnxType.TypeVariable("T", List.of(OnnxType.tensor(OnnxType.float32()), OnnxType.tensor(OnnxType.float16())))),
757 ;
758
759 final OnnxType.TypeVariable typeVariable;
760
761 TypeConstraint(OnnxType.TypeVariable typeVariable) {
762 assert typeVariable.name().equals(name());
763 this.typeVariable = typeVariable;
764 }
765
766 @Override
767 public OnnxType.TypeVariable typeVariable() {
768 return typeVariable;
769 }
770 }
771
772 public enum InputParameter implements OnnxParameter {
773 input(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
774 skip(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
775 gamma(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
776 bias(TypeConstraint.T.typeVariable(), Quantifier.OPTIONAL),
777 ;
778
779 final OnnxType type;
780 final Quantifier quantifier;
781
782 InputParameter(OnnxType type, Quantifier quantifier) {
783 this.type = type;
784 this.quantifier = quantifier;
785 }
786
787 @Override
788 public OnnxType type() {
789 return type;
790 }
791
792 @Override
793 public Quantifier quantifier() {
794 return quantifier;
795 }
796 }
797
798 public enum OutputParameter implements OnnxParameter {
799 output(TypeConstraint.T.typeVariable(), Quantifier.REQUIRED),
800 mean(OnnxType.TENSOR_FLOAT32, Quantifier.OPTIONAL),
801 inv_std_var(OnnxType.TENSOR_FLOAT32, Quantifier.OPTIONAL),
802 input_skip_bias_sum(OnnxType.TENSOR_FLOAT32, Quantifier.OPTIONAL),
803 ;
804
805 final OnnxType type;
806 final Quantifier quantifier;
807
808 OutputParameter(OnnxType type, Quantifier quantifier) {
809 this.type = type;
810 this.quantifier = quantifier;
811 }
812
813 @Override
814 public OnnxType type() {
815 return type;
816 }
817
818 @Override
819 public Quantifier quantifier() {
820 return quantifier;
821 }
822 }
823
824 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
825 NAME,
826 List.of(Attribute.values()),
827 List.of(TypeConstraint.values()),
828 List.of(InputParameter.values()),
829 List.of(OutputParameter.values())
830 );
831
832 public SkipSimplifiedLayerNormalization(ExternalizedOp def) {
833 super(SCHEMA, def);
834 }
835
836 SkipSimplifiedLayerNormalization(SkipSimplifiedLayerNormalization that, CodeContext cc) {
837 super(that, cc);
838 }
839
840 @Override
841 public SkipSimplifiedLayerNormalization transform(CodeContext cc, CodeTransformer ot) {
842 return new SkipSimplifiedLayerNormalization(this, cc);
843 }
844
845 SkipSimplifiedLayerNormalization(CodeType resultType, Set<OutputParameter> optionalOutputs, Value input, Value skip, Value gamma, java.util.Optional<Value> bias, java.util.Optional<Float> epsilon) {
846 super(SCHEMA, resultType, optionalOutputs, List.of(input, skip, gamma, bias), List.of(epsilon));
847 }
848
849 @Override
850 public SequencedSet<OnnxParameter> onnxOutputs() {
851 return onnxOutputs(SCHEMA);
852 }
853
854 @Override
855 public SequencedMap<OnnxParameter, Object> onnxInputs() {
856 return onnxInputs(SCHEMA, List.of(input(), skip(), gamma(), bias()));
857 }
858
859 public Value input() {
860 return operands().get(0);
861 }
862
863 public Value skip() {
864 return operands().get(1);
865 }
866
867 public Value gamma() {
868 return operands().get(2);
869 }
870
871 public java.util.Optional<Value> bias() {
872 int i = optionalInputArguments.indexOf(InputParameter.bias);
873 return i != -1 ? java.util.Optional.of(operands().get(3 + i)) : java.util.Optional.empty();
874 }
875 }
876
877 public static SkipSimplifiedLayerNormalization SkipSimplifiedLayerNormalization(CodeType resultType, Set<SkipSimplifiedLayerNormalization.OutputParameter> optionalOutputs, Value input, Value skip, Value gamma, java.util.Optional<Value> bias, java.util.Optional<Float> epsilon) {
878 return new SkipSimplifiedLayerNormalization(resultType, optionalOutputs, input, skip, gamma, bias, epsilon);
879 }
880
881 // @@@ this should be generated from onnxruntime-extensions
882 @OpFactoryHelper.OpDeclaration(CLIPTokenizer.NAME)
883 public static final class CLIPTokenizer extends OnnxOp {
884 public static final String NAME = "ai.onnx.contrib.CLIPTokenizer";
885
886 public enum Attribute implements OnnxAttribute {
887 vocab(String.class, false, null),
888 merges(String.class, false, null),
889 padding_length(Long.class, true, -1),
890 ;
891
892 final Class<?> t;
893 final boolean optional;
894 final Object defaultValue;
895
896 Attribute(Class<?> type, boolean optional, Object defaultValue) {
897 this.t = type;
898 this.optional = optional;
899 this.defaultValue = defaultValue;
900 assert optional || defaultValue == null;
901 }
902
903 public Class<?> type() {
904 return t;
905 }
906
907 public boolean isOptional() {
908 return optional;
909 }
910
911 public Object defaultValue() {
912 return defaultValue;
913 }
914 }
915
916 public enum TypeConstraint implements OnnxTypeConstraint.None { }
917
918 public enum InputParameter implements OnnxParameter {
919 input_text(OnnxType.TENSOR_STRING, Quantifier.REQUIRED),
920 ;
921
922 final OnnxType type;
923 final Quantifier quantifier;
924
925 InputParameter(OnnxType type, Quantifier quantifier) {
926 this.type = type;
927 this.quantifier = quantifier;
928 }
929
930 @Override
931 public OnnxType type() {
932 return type;
933 }
934
935 @Override
936 public Quantifier quantifier() {
937 return quantifier;
938 }
939 }
940
941 public enum OutputParameter implements OnnxParameter {
942 input_ids(OnnxType.TENSOR_INT64, Quantifier.REQUIRED),
943 attention_mask(OnnxType.TENSOR_INT64, Quantifier.OPTIONAL),
944 offset_mapping(OnnxType.TENSOR_INT64, Quantifier.OPTIONAL),
945 ;
946
947 final OnnxType type;
948 final Quantifier quantifier;
949
950 OutputParameter(OnnxType type, Quantifier quantifier) {
951 this.type = type;
952 this.quantifier = quantifier;
953 }
954
955 @Override
956 public OnnxType type() {
957 return type;
958 }
959
960 @Override
961 public Quantifier quantifier() {
962 return quantifier;
963 }
964 }
965
966 public static final OnnxSchema SCHEMA = new OnnxSchemaRecord(
967 NAME,
968 List.of(Attribute.values()),
969 List.of(TypeConstraint.values()),
970 List.of(InputParameter.values()),
971 List.of(OutputParameter.values())
972 );
973
974 public CLIPTokenizer(ExternalizedOp def) {
975 super(SCHEMA, def);
976 }
977
978 CLIPTokenizer(CLIPTokenizer that, CodeContext cc) {
979 super(that, cc);
980 }
981
982 @Override
983 public CLIPTokenizer transform(CodeContext cc, CodeTransformer ot) {
984 return new CLIPTokenizer(this, cc);
985 }
986
987 CLIPTokenizer(CodeType resultType, Set<OutputParameter> optionalOutputs, Value input_text, String vocab, String merges, java.util.Optional<Long> padding_length) {
988 super(SCHEMA, resultType, optionalOutputs, List.of(input_text), List.of(vocab, merges, padding_length));
989 }
990
991 @Override
992 public SequencedSet<OnnxParameter> onnxOutputs() {
993 return onnxOutputs(SCHEMA);
994 }
995
996 @Override
997 public SequencedMap<OnnxParameter, Object> onnxInputs() {
998 return onnxInputs(SCHEMA, List.of(input_text()));
999 }
1000
1001 public Value input_text() {
1002 return operands().get(0);
1003 }
1004 }
1005
1006 public static CLIPTokenizer CLIPTokenizer(CodeType resultType, Set<CLIPTokenizer.OutputParameter> optionalOutputs, Value input_text, String vocab, String merges, java.util.Optional<Long> padding_length) {
1007 return new CLIPTokenizer(resultType, optionalOutputs, input_text, vocab, merges, padding_length);
1008 }
1009
1010
1011 @OpFactoryHelper.OpDeclaration(If.NAME)
1012 public static final class If extends OnnxOp implements Nested {
1013 public static final String NAME = "If";
1014
1015 final Body thenBody, elseBody;
1016
1017 // @@@ make or fake elseBody as "else_branch" attribute and thenBody as "then_branch" attribute
1018 public enum Attribute implements OnnxOp.OnnxAttribute.None { }
1019
1020 public enum TypeConstraint implements OnnxOp.OnnxTypeConstraint {
1021 V(new OnnxType.TypeVariable("V", List.of(OnnxType.tensor(OnnxType.uint8()), OnnxType.tensor(OnnxType.uint16()), OnnxType.tensor(OnnxType.uint32()), OnnxType.tensor(OnnxType.uint64()), OnnxType.tensor(OnnxType.int8()), OnnxType.tensor(OnnxType.int16()), OnnxType.tensor(OnnxType.int32()), OnnxType.tensor(OnnxType.int64()), OnnxType.tensor(OnnxType.bfloat16()), OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.float32()), OnnxType.tensor(OnnxType.float64()), OnnxType.tensor(OnnxType.bool())))),
1022 B(new OnnxType.TypeVariable("B", List.of(OnnxType.tensor(OnnxType.bool())))),
1023 ;
1024
1025 final OnnxType.TypeVariable typeVariable;
1026
1027 TypeConstraint(OnnxType.TypeVariable typeVariable) {
1028 assert typeVariable.name().equals(name());
1029 this.typeVariable = typeVariable;
1030 }
1031
1032 @Override
1033 public OnnxType.TypeVariable typeVariable() {
1034 return typeVariable;
1035 }
1036 }
1037
1038 public enum InputParameter implements OnnxOp.OnnxParameter {
1039 cond(TypeConstraint.B.typeVariable(), OnnxOp.OnnxParameter.Quantifier.REQUIRED),
1040 ;
1041
1042 final OnnxType type;
1043 final OnnxOp.OnnxParameter.Quantifier quantifier;
1044
1045 InputParameter(OnnxType type, OnnxOp.OnnxParameter.Quantifier quantifier) {
1046 this.type = type;
1047 this.quantifier = quantifier;
1048 }
1049
1050 @Override
1051 public OnnxType type() {
1052 return type;
1053 }
1054
1055 @Override
1056 public OnnxOp.OnnxParameter.Quantifier quantifier() {
1057 return quantifier;
1058 }
1059 }
1060
1061 public enum OutputParameter implements OnnxOp.OnnxParameter {
1062 output(TypeConstraint.V.typeVariable(), OnnxOp.OnnxParameter.Quantifier.VARIADIC),
1063 ;
1064
1065 final OnnxType type;
1066 final OnnxOp.OnnxParameter.Quantifier quantifier;
1067
1068 OutputParameter(OnnxType type, OnnxOp.OnnxParameter.Quantifier quantifier) {
1069 this.type = type;
1070 this.quantifier = quantifier;
1071 }
1072
1073 @Override
1074 public OnnxType type() {
1075 return type;
1076 }
1077
1078 @Override
1079 public OnnxOp.OnnxParameter.Quantifier quantifier() {
1080 return quantifier;
1081 }
1082 }
1083
1084 public static final OnnxOp.OnnxSchema SCHEMA = new OnnxSchemaRecord(
1085 NAME,
1086 List.of(Attribute.values()),
1087 List.of(TypeConstraint.values()),
1088 List.of(InputParameter.values()),
1089 List.of(OutputParameter.values())
1090 );
1091
1092 public If(ExternalizedOp def) {
1093 super(SCHEMA, def);
1094
1095 this.thenBody = def.bodyDefinitions().get(0).build(this);
1096 this.elseBody = def.bodyDefinitions().get(1).build(this);
1097 }
1098
1099 If(If that, CodeContext cc, CodeTransformer ot) {
1100 super(that, cc);
1101
1102 this.thenBody = that.thenBody.transform(cc, ot).build(this);
1103 this.elseBody = that.elseBody.transform(cc, ot).build(this);
1104 }
1105
1106 @Override
1107 public If transform(CodeContext cc, CodeTransformer ot) {
1108 return new If(this, cc, ot);
1109 }
1110
1111 If(CodeType resultType, Value cond, Body.Builder thenBranch, Body.Builder elseBranch) {
1112 super(SCHEMA, resultType, Set.of(), List.of(cond), List.of());
1113
1114 this.thenBody = thenBranch.build(this);
1115 this.elseBody = elseBranch.build(this);
1116 }
1117
1118 @Override
1119 public List<Body> bodies() {
1120 return List.of(thenBody, elseBody);
1121 }
1122
1123 @Override
1124 public SequencedSet<OnnxOp.OnnxParameter> onnxOutputs() {
1125 return onnxOutputs(SCHEMA);
1126 }
1127
1128 @Override
1129 public SequencedMap<OnnxOp.OnnxParameter, Object> onnxInputs() {
1130 return onnxInputs(SCHEMA, List.of(cond()));
1131 }
1132
1133 public Value cond() {
1134 return operands().get(0);
1135 }
1136
1137 public Body elseBranch() {
1138 return elseBody;
1139 }
1140
1141 public Body thenBranch() {
1142 return thenBody;
1143 }
1144 }
1145
1146 public static If If(CodeType resultType, Value cond, Body.Builder thenBody, Body.Builder elseBody) {
1147 return new If(resultType, cond, thenBody, elseBody);
1148 }
1149
1150 @OpFactoryHelper.OpDeclaration(Loop.NAME)
1151 public static final class Loop extends OnnxOp implements Op.Loop {
1152 public static final String NAME = "Loop";
1153
1154 final Body body;
1155
1156 // @@@ make or fake body
1157 public enum Attribute implements OnnxOp.OnnxAttribute.None { }
1158
1159 public enum TypeConstraint implements OnnxOp.OnnxTypeConstraint {
1160 V(new OnnxType.TypeVariable("V", List.of(OnnxType.tensor(OnnxType.uint8()), OnnxType.tensor(OnnxType.uint16()), OnnxType.tensor(OnnxType.uint32()), OnnxType.tensor(OnnxType.uint64()), OnnxType.tensor(OnnxType.int8()), OnnxType.tensor(OnnxType.int16()), OnnxType.tensor(OnnxType.int32()), OnnxType.tensor(OnnxType.int64()), OnnxType.tensor(OnnxType.bfloat16()), OnnxType.tensor(OnnxType.float16()), OnnxType.tensor(OnnxType.float32()), OnnxType.tensor(OnnxType.float64()), OnnxType.tensor(OnnxType.bool())))),
1161 I(new OnnxType.TypeVariable("I", List.of(OnnxType.tensor(OnnxType.int64())))),
1162 B(new OnnxType.TypeVariable("B", List.of(OnnxType.tensor(OnnxType.bool())))),
1163 ;
1164
1165 final OnnxType.TypeVariable typeVariable;
1166
1167 TypeConstraint(OnnxType.TypeVariable typeVariable) {
1168 assert typeVariable.name().equals(name());
1169 this.typeVariable = typeVariable;
1170 }
1171
1172 @Override
1173 public OnnxType.TypeVariable typeVariable() {
1174 return typeVariable;
1175 }
1176 }
1177
1178 public enum InputParameter implements OnnxOp.OnnxParameter {
1179 // @@@ Onnx spec declares the input parameters as optional, however it is causing problems
1180 M(TypeConstraint.I.typeVariable(), OnnxOp.OnnxParameter.Quantifier.REQUIRED),
1181 cond(TypeConstraint.B.typeVariable(), OnnxOp.OnnxParameter.Quantifier.REQUIRED),
1182 v_initial(TypeConstraint.V.typeVariable(), OnnxOp.OnnxParameter.Quantifier.VARIADIC),
1183 ;
1184
1185 final OnnxType type;
1186 final OnnxOp.OnnxParameter.Quantifier quantifier;
1187
1188 InputParameter(OnnxType type, OnnxOp.OnnxParameter.Quantifier quantifier) {
1189 this.type = type;
1190 this.quantifier = quantifier;
1191 }
1192
1193 @Override
1194 public OnnxType type() {
1195 return type;
1196 }
1197
1198 @Override
1199 public OnnxOp.OnnxParameter.Quantifier quantifier() {
1200 return quantifier;
1201 }
1202 }
1203
1204 public enum OutputParameter implements OnnxOp.OnnxParameter {
1205 v_final_and_scan_outputs(TypeConstraint.V.typeVariable(), OnnxOp.OnnxParameter.Quantifier.VARIADIC),
1206 ;
1207
1208 final OnnxType type;
1209 final OnnxOp.OnnxParameter.Quantifier quantifier;
1210
1211 OutputParameter(OnnxType type, OnnxOp.OnnxParameter.Quantifier quantifier) {
1212 this.type = type;
1213 this.quantifier = quantifier;
1214 }
1215
1216 @Override
1217 public OnnxType type() {
1218 return type;
1219 }
1220
1221 @Override
1222 public OnnxOp.OnnxParameter.Quantifier quantifier() {
1223 return quantifier;
1224 }
1225 }
1226
1227 public static final OnnxOp.OnnxSchema SCHEMA = new OnnxSchemaRecord(
1228 NAME,
1229 List.of(Attribute.values()),
1230 List.of(TypeConstraint.values()),
1231 List.of(InputParameter.values()),
1232 List.of(OutputParameter.values())
1233 );
1234
1235 public Loop(ExternalizedOp def) {
1236 super(SCHEMA, def);
1237
1238 this.body = def.bodyDefinitions().get(0).build(this);
1239 }
1240
1241 Loop(ExplicitOnnxOps.Loop that, CodeContext cc, CodeTransformer ot) {
1242 super(that, cc);
1243
1244 this.body = that.body.transform(cc, ot).build(this);
1245 }
1246
1247 @Override
1248 public ExplicitOnnxOps.Loop transform(CodeContext cc, CodeTransformer ot) {
1249 return new ExplicitOnnxOps.Loop(this, cc, ot);
1250 }
1251
1252 Loop(CodeType resultType, Value m, Value cond, Object v_initial, Body.Builder body) {
1253 super(SCHEMA, resultType, Set.of(), List.of(m, cond, v_initial), List.of());
1254
1255 this.body = body.build(this);
1256 }
1257
1258 @Override
1259 public List<Body> bodies() {
1260 return List.of(body);
1261 }
1262
1263 @Override
1264 public SequencedSet<OnnxOp.OnnxParameter> onnxOutputs() {
1265 return onnxOutputs(SCHEMA);
1266 }
1267
1268 @Override
1269 public SequencedMap<OnnxOp.OnnxParameter, Object> onnxInputs() {
1270 return onnxInputs(SCHEMA, List.of(cond()));
1271 }
1272
1273 public Value max() {
1274 return operands().get(0);
1275 }
1276
1277 public Value cond() {
1278 return operands().get(1);
1279 }
1280
1281 public List<Value> v_initial() {
1282 return operands().subList(2, operands().size());
1283 }
1284
1285 @Override
1286 public Body loopBody() {
1287 return body;
1288 }
1289 }
1290
1291 public static Loop Loop(CodeType resultType, Value m, Value cond, Object v_initial, Body.Builder body) {
1292 return new Loop(resultType, m, cond, v_initial, body);
1293 }
1294 }