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 jdk.incubator.code.*;
29 import jdk.incubator.code.extern.ExternalizedOp;
30
31 import java.util.*;
32
33 public abstract class OnnxOp extends AbstractOp implements ExternalizedOp.Externalizable {
34
35 public interface OnnxAttribute {
36 String name();
37
38 Class<?> type();
39
40 Object defaultValue();
41
42 boolean isOptional();
43
44 default void process(Map<String, Object> attrs, Object value) {
45 if (value instanceof Optional<?> o) {
46 value = o.orElse(null);
47 }
48 // @@@ Parse attribute from string value
49 // @@@ Arrays don't serialize
50 if (type().isInstance(value)) {
51 attrs.put(name(), value);
52 } else if (value == null) {
53 // Ignore
54 } else {
55 throw new UnsupportedOperationException();
56 }
57 }
58
59 default <T> T access(Class<T> type, Map<String, Object> attrs) {
60 Object value = attrs.get(name());
61 if (value == null && !isOptional()) {
62 throw new NoSuchElementException();
63 }
64 return type.cast(value);
65 }
66
67 static Map<String, Object> process(ExternalizedOp eop,
68 List<OnnxAttribute> attributes) {
69 Map<String, Object> attrs = new HashMap<>();
70 for (OnnxAttribute attribute : attributes) {
71 Object v = eop.attributes().get(attribute.name());
72 if (v == null && !attribute.isOptional()) {
73 throw new NoSuchElementException(attribute.name());
74 }
75 attribute.process(attrs, v);
76 }
77
78 return Map.copyOf(attrs);
79 }
80
81 interface None extends OnnxAttribute {
82 @Override
83 default String name() {
84 throw new UnsupportedOperationException();
85 }
86
87 @Override
88 default Class<?> type() {
89 throw new UnsupportedOperationException();
90 }
91
92 @Override
93 default Object defaultValue() {
94 throw new UnsupportedOperationException();
95 }
96
97 @Override
98 default boolean isOptional() {
99 throw new UnsupportedOperationException();
100 }
101 }
102
103 }
104
105 public interface OnnxTypeConstraint {
106 String name();
107
108 OnnxType.TypeVariable typeVariable();
109
110 interface None extends OnnxTypeConstraint {
111 @Override
112 default String name() {
113 throw new UnsupportedOperationException();
114 }
115
116 @Override
117 default OnnxType.TypeVariable typeVariable() {
118 throw new UnsupportedOperationException();
119 }
120 }
121 }
122
123 public interface OnnxParameter {
124 enum Quantifier {
125 REQUIRED, // Exactly once
126 OPTIONAL, // Once or none
127 VARIADIC, // One or more
128 ;
129
130 public boolean isOptional() {
131 return this == OPTIONAL;
132 }
133
134 public boolean isRequired() {
135 return this == REQUIRED;
136 }
137
138 public boolean isVariadoc() {
139 return this == VARIADIC;
140 }
141 }
142
143 String name();
144
145 OnnxType type();
146
147 Quantifier quantifier();
148
149 interface None extends OnnxParameter {
150 @Override
151 default String name() {
152 throw new UnsupportedOperationException();
153 }
154
155 @Override
156 default OnnxType type() {
157 throw new UnsupportedOperationException();
158 }
159
160 @Override
161 default Quantifier quantifier() {
162 throw new UnsupportedOperationException();
163 }
164 }
165 }
166
167 public interface OnnxSchema {
168 String name();
169
170 List<OnnxAttribute> attributes();
171
172 List<OnnxTypeConstraint> typeConstraints();
173
174 List<OnnxParameter> inputs();
175
176 List<OnnxParameter> outputs();
177 }
178
179 record OnnxSchemaRecord(
180 String name,
181 List<OnnxAttribute> attributes,
182 List<OnnxTypeConstraint> typeConstraints,
183 List<OnnxParameter> inputs,
184 List<OnnxParameter> outputs
185 ) implements OnnxSchema {}
186
187 static List<Value> concatValues(Value operand) {
188 return List.of(operand);
189 }
190
191 static List<Value> concatValues(Value... operands) {
192 return List.of(operands);
193 }
194
195 static List<Value> concatValues(List<Object> operands) {
196 return concatValues(operands.toArray());
197 }
198
199 static List<Value> concatValues(Object... operands) {
200 List<Value> l = new ArrayList<>();
201 for (Object operand : operands) {
202 switch (operand) {
203 case Value v -> l.add(v);
204 case Optional<?> ov -> {
205 ov.ifPresent(o -> l.add((Value) o));
206 }
207 case List<?> vs -> {
208 for (Object v : vs) {
209 l.add((Value) v);
210 }
211 }
212 default -> throw new UnsupportedOperationException();
213 }
214 }
215 return l;
216 }
217
218 static final String ATTRIBUTE_OPTIONAL_INPUTS = "optional_inputs";
219 static final String ATTRIBUTE_OPTIONAL_OUTPUTS = "optional_outputs";
220
221 final OnnxSchema schema;
222 final Map<String, Object> onnxAttributes;
223 final CodeType resultType;
224 final List<OnnxParameter> optionalInputArguments;
225 final List<OnnxParameter> optionalOutputParameters;
226
227 @SuppressWarnings("unchecked")
228 OnnxOp(OnnxSchema schema, ExternalizedOp def) {
229 super(def.operands());
230
231 this.schema = schema;
232 this.onnxAttributes = schema.attributes().isEmpty()
233 ? Map.of()
234 : OnnxAttribute.process(def, schema.attributes());
235 this.resultType = def.resultType();
236
237 // @@@ Filter optional
238 this.optionalInputArguments = switch (def.attributes().get(ATTRIBUTE_OPTIONAL_INPUTS)) {
239 case List<?> s -> (List<OnnxParameter>) s;
240 case null -> List.of();
241 default -> throw new UnsupportedOperationException();
242 };
243
244 // @@@ Filter optional
245 this.optionalOutputParameters = switch (def.attributes().get(ATTRIBUTE_OPTIONAL_OUTPUTS)) {
246 case List<?> s -> (List<OnnxParameter>) s;
247 case null -> List.of();
248 default -> throw new UnsupportedOperationException();
249 };
250 }
251
252 OnnxOp(OnnxOp that, CodeContext cc) {
253 super(that, cc);
254
255 this.schema = that.schema;
256 this.onnxAttributes = Map.copyOf(that.onnxAttributes);
257 this.resultType = that.resultType;
258 this.optionalInputArguments = List.copyOf(that.optionalInputArguments);
259 this.optionalOutputParameters = List.copyOf(that.optionalOutputParameters);
260 }
261
262 OnnxOp(OnnxSchema schema, CodeType resultType,
263 Set<? extends OnnxParameter> optionalOutputParameters,
264 List<Object> inputArguments,
265 List<Object> attributeValues) {
266 super(concatValues(inputArguments));
267
268 this.schema = schema;
269 this.resultType = resultType;
270
271 // Optional output parameters
272
273 if (!optionalOutputParameters.isEmpty()) {
274 List<OnnxParameter> l = new ArrayList<>();
275
276 for (int i = 0; i < schema.outputs().size(); i++) {
277 OnnxParameter p = schema.outputs().get(i);
278 if (p.quantifier().isOptional()
279 && optionalOutputParameters.contains(p)) {
280 l.add(p);
281 }
282 }
283 this.optionalOutputParameters = List.copyOf(l);
284 } else {
285 this.optionalOutputParameters = List.of();
286 }
287
288 // Optional input parameters
289
290 if (!inputArguments.isEmpty()) {
291 List<OnnxParameter> l = new ArrayList<>();
292
293 for (int i = 0; i < schema.inputs().size(); i++) {
294 OnnxParameter p = schema.inputs().get(i);
295 if (p.quantifier().isOptional()) {
296 assert inputArguments.get(i) instanceof Optional;
297 if (inputArguments.get(i) instanceof Optional<?> optionalValue
298 && optionalValue.isPresent()) {
299 l.add(p);
300 }
301 }
302 }
303 if (!l.isEmpty()) {
304 this.optionalInputArguments = List.copyOf(l);
305 } else {
306 this.optionalInputArguments = List.of();
307 }
308 } else {
309 this.optionalInputArguments = List.of();
310 }
311
312 // Attributes
313
314 if (!attributeValues.isEmpty()) {
315 Map<String, Object> attrs = new HashMap<>();
316 assert schema.attributes().size() == attributeValues.size();
317 for (int i = 0; i < schema.attributes().size(); i++) {
318 schema.attributes().get(i).process(attrs, attributeValues.get(i));
319 }
320 this.onnxAttributes = Map.copyOf(attrs);
321 } else {
322 this.onnxAttributes = Map.of();
323 }
324 }
325
326 @Override
327 public CodeType resultType() {
328 return resultType;
329 }
330
331 @Override
332 public String externalizeOpName() {
333 return schema.name();
334 }
335
336 @Override
337 public Map<String, Object> externalize() {
338 HashMap<String, Object> m = new HashMap<>(onnxAttributes);
339 if (!optionalInputArguments.isEmpty()) {
340 m.put(ATTRIBUTE_OPTIONAL_INPUTS, optionalInputArguments);
341 }
342 if (!optionalOutputParameters.isEmpty()) {
343 m.put(ATTRIBUTE_OPTIONAL_OUTPUTS, optionalOutputParameters);
344 }
345 return Collections.unmodifiableMap(m);
346 }
347
348 public OnnxSchema schema() {
349 return schema;
350 }
351
352 // @@@ Change to Map<OnnxAttribute, Object>
353 public Map<String, Object> onnxAttributes() {
354 return onnxAttributes;
355 }
356
357 public SequencedSet<OnnxParameter> onnxOutputs() {
358 return Collections.emptyNavigableSet();
359 }
360
361 SequencedSet<OnnxParameter> onnxOutputs(OnnxSchema schema) {
362 LinkedHashSet<OnnxParameter> s = new LinkedHashSet<>();
363 for (OnnxParameter p : schema.outputs()) {
364 if (!p.quantifier().isOptional() || optionalOutputParameters.contains(p)) {
365 s.add(p);
366 }
367 }
368
369 return s;
370 }
371
372 public SequencedMap<OnnxParameter, Object> onnxInputs() {
373 return Collections.emptyNavigableMap();
374 }
375
376 SequencedMap<OnnxParameter, Object> onnxInputs(OnnxSchema schema, List<Object> inputArguments) {
377 assert schema.inputs().size() == inputArguments.size();
378 if (!inputArguments.isEmpty()) {
379 SequencedMap<OnnxParameter, Object> inputs = new LinkedHashMap<>();
380 for (int i = 0; i < schema.inputs().size(); i++) {
381 inputs.put(schema.inputs().get(i), inputArguments.get(i));
382 }
383 return inputs;
384 } else {
385 return Collections.emptyNavigableMap();
386 }
387 }
388 }