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;
27
28 import java.lang.foreign.ValueLayout;
29 import java.util.*;
30 import java.util.function.Function;
31 import java.util.stream.IntStream;
32 import jdk.incubator.code.Block;
33 import jdk.incubator.code.CodeItem;
34 import jdk.incubator.code.Op;
35 import jdk.incubator.code.Value;
36 import jdk.incubator.code.dialect.core.CoreOp;
37 import jdk.incubator.code.dialect.core.TupleType;
38 import jdk.incubator.code.dialect.java.JavaOp;
39 import jdk.incubator.code.dialect.java.JavaType;
40 import jdk.incubator.code.extern.OpWriter;
41 import oracle.code.onnx.ir.OnnxOp;
42 import oracle.code.onnx.ir.OnnxOps;
43 import oracle.code.onnx.ir.OnnxType;
44 import oracle.code.onnx.proto.OnnxBuilder.*;
45 import oracle.code.onnx.proto.OnnxConstants.*;
46
47 public final class OnnxProtoBuilder {
48
49 static final int IR_VERSION = 10;
50 static final int OPSET_VERSION = 21;
51
52 private static final class Indexer {
53
54 private final Function<CodeItem, String> baseNames;
55 private final HashMap<String, String> remap;
56
57
58 Indexer(Op root, Map<Value, String> explicitNames) {
59 this.baseNames = OpWriter.computeGlobalNames(root);
60 this.remap = new HashMap<>();
61 explicitNames.forEach(this::setName);
62 }
63
64 void setName(Value val, String name) {
65 switch (val) {
66 case Op.Result or when or.op() instanceof CoreOp.TupleOp to -> {
67 remap.put(baseName(val), name);
68 for (int i = 0; i < to.operands().size(); i++) {
69 setName(to.operands().get(i), name + "." + i);
70 }
71 }
72 case Block.Parameter _ when val.type() instanceof TupleType tt -> {
73 for (int i = 0; i < tt.componentTypes().size(); i++) {
74 remap.put(baseName(val, i), name +"." + i);
75 }
76 }
77 default -> {
78 remap.put(baseName(val), name);
79 if (val instanceof Op.Result or && or.op() instanceof CoreOp.TupleLoadOp tlo) {
80 Value tr = tlo.operands().getFirst();
81 remap.put(baseName(tr, tlo.index()), name);
82 if (tr instanceof Op.Result tor && tor.op() instanceof CoreOp.TupleOp to) {
83 setName(to.operands().get(tlo.index()), name);
84 }
85 }
86 }
87 }
88 }
89
90 private String baseName(Value value) {
91 return "%" + baseNames.apply(value);
92 }
93
94 private String baseName(Value value, int elementIndex) {
95 var name = baseName(value);
96 return elementIndex > 0 ? name + '.' + elementIndex : name;
97 }
98
99 String nameOf(Value value) {
100 var name = baseName(value);
101 return remap.getOrDefault(name, name);
102 }
103
104 String nameOf(Value tuple, int elementIndex) {
105 var name = baseName(tuple, elementIndex);
106 return remap.getOrDefault(name, name);
107 }
108
109 void mapTupleLoad(Value tupleLoadResult, Value tuple, int elementIndex) {
110 remap.putIfAbsent(baseName(tupleLoadResult), nameOf(tuple, elementIndex));
111 }
112
113 void mapTupleElements(Value tuple, List<Value> elements) {
114 for (int i = 0; i < elements.size(); i++) {
115 remap.putIfAbsent(baseName(tuple, i), nameOf(elements.get(i)));
116 }
117 }
118 }
119
120 public static byte[] buildModel(String domain, CoreOp.ModuleOp module, List<Object> initializers) {
121 return buildModel(domain, module, initializers, Map.of(), _ -> null);
122 }
123
124 public record ExternalTensorDataInfo(String location, long offset, long length) {
125 }
126
127 public static byte[] buildModel(String domain, CoreOp.ModuleOp module, List<Object> initializers, Map<Value, String> explicitValueNames, Function<Tensor, ExternalTensorDataInfo> tensorDataExternalizer) {
128 var indexer = new Indexer(module, explicitValueNames);
129
130 var functions = new ArrayList<>(module.functionTable().sequencedValues());
131 var imports = new ArrayList<String>();
132 if (functions.size() > 1) imports.add(domain); // self domain import if additional functions
133 for (var f : functions) {
134 for (var op : f.body().entryBlock().ops()) { // auto import of op domains
135 if (op instanceof OnnxOp oop) {
136 String name = oop.schema().name();
137 int di = name.lastIndexOf('.');
138 if (di > 0) {
139 String dn = name.substring(0, di);
140 if (!imports.contains(dn)) imports.add(dn);
141 }
142 }
143 }
144 }
145 var mainFunc = functions.removeLast();
146 var mainBlock = mainFunc.body().entryBlock();
147
148 return buildModel(
149 graph(domain, mainFunc.funcName(), indexer, mainBlock, initializers, 0, tensorDataExternalizer),
150 imports,
151 functions.stream().map(f ->
152 function(domain, imports, f.funcName(),
153 expandTuples(indexer, f.parameters()),
154 expandTuples(indexer, f.body().entryBlock().terminatingOp().operands()),
155 nodes(domain, indexer, f.body().entryBlock().ops()))).toList());
156 }
157
158 // @@@ unchecked constraints:
159 // tensor FuncOp parameters and single tensor return type
160 // OnnxOps (with tensor operands and single tensor return value) and ReturnOp (returning single tensor)
161 // entry block only
162 static byte[] buildModel(Block block, List<Tensor> initializers) {
163 var indexer = new Indexer(block.ancestorOp(), Map.of());
164 var model = buildModel(graph(null, null, indexer, block, initializers, 0), List.of(), List.of());
165 // System.out.println(OnnxModel.readFrom(model).toText());
166 return model;
167 }
168
169 static byte[] buildModel(List<TensorProto> initializers, List<ValueInfoProto> inputs, List<NodeProto> ops, List<String> outputNames) {
170 return buildModel(graph(null, initializers, inputs, ops, outputNames), List.of(), List.of());
171 }
172
173 static byte[] buildModel(List<TensorProto> initializers, List<ValueInfoProto> inputs, List<NodeProto> ops, List<String> outputNames, List<String> customImportDomains, List<FunctionProto> functions) {
174 return buildModel(graph(null, initializers, inputs, ops, outputNames), customImportDomains, functions);
175 }
176
177 static byte[] buildModel(GraphProto graph, List<String> imports, List<FunctionProto> functions) {
178 return new ModelProto()
179 .irVersion(IR_VERSION)
180 .opsetImport(new OperatorSetIdProto().version(OPSET_VERSION))
181 .forEach(imports, (m, d) -> m.opsetImport(new OperatorSetIdProto().domain(d).version(1)))
182 .forEach(functions, ModelProto::functions)
183 .graph(graph)
184 .getBytes();
185 }
186
187 static List<String> expandTuples(Indexer indexer, List<? extends Value> values) {
188 var names = new ArrayList<String>();
189 expandTuples(indexer, names, values);
190 return names;
191 }
192
193 static void expandTuples(Indexer indexer, List<String> names, List<? extends Value> values) {
194 for (var v : values) {
195 if (v instanceof Op.Result or && or.op() instanceof CoreOp.TupleOp op) {
196 expandTuples(indexer, names, op.operands());
197 } else if (v instanceof Op.Result or && or.op() instanceof CoreOp.TupleLoadOp op) {
198 names.add(indexer.nameOf(op.operands().getFirst(), op.index()));
199 } else if (v.type() instanceof TupleType tt) {
200 var ct = tt.componentTypes();
201 for (int i = 0; i < ct.size(); i++) {
202 names.add(indexer.nameOf(v, i));
203 }
204 } else {
205 names.add(indexer.nameOf(v));
206 }
207 }
208 }
209
210 static GraphProto graph(String domain, String graphName, Indexer indexer, Block block, List<?> initializers, int scalarArgs) {
211 return graph(domain, graphName, indexer, block, initializers, scalarArgs, _ -> null);
212 }
213
214 static GraphProto graph(String domain, String graphName, Indexer indexer, Block block, List<?> initializers, int scalarArgs, Function<Tensor, ExternalTensorDataInfo> tensorDataExternalizer) {
215 var params = block.parameters();
216 params.forEach(indexer::nameOf);
217 int firstInitializer = params.size() - initializers.size();
218 var args = params.subList(0, firstInitializer);
219 return graph(graphName,
220 IntStream.range(0, initializers.size()).boxed().<TensorProto>mapMulti((i, tps) -> {
221 Object val = initializers.get(i);
222 if (val instanceof Record) {
223 var rcs = val.getClass().getRecordComponents();
224 for (int rci = 0; rci < rcs.length; rci++) try {
225 tps.accept(tensorProto(indexer.nameOf(params.get(i + firstInitializer), rci), (Tensor)(rcs[rci].getAccessor().invoke(val)), tensorDataExternalizer));
226 } catch (ReflectiveOperationException e) {
227 throw new IllegalArgumentException(e);
228 }
229 } else if (val instanceof Tensor[] tarr) {
230 for (int tai = 0; tai < tarr.length; tai++) {
231 tps.accept(tensorProto(indexer.nameOf(params.get(i + firstInitializer), tai), tarr[tai], tensorDataExternalizer));
232 }
233 } else {
234 tps.accept(tensorProto(indexer.nameOf(params.get(i + firstInitializer)), (Tensor)val, tensorDataExternalizer));
235 }
236 }).toList(),
237 tensorInfos(indexer, args, scalarArgs),
238 nodes(domain, indexer, block.ops()),
239 expandTuples(indexer, block.terminatingOp().operands()));
240 }
241
242 static List<String> opInputNames(Indexer indexer, SequencedMap<OnnxOp.OnnxParameter, Object> inputs) {
243 List<String> inputNames = inputs.sequencedValues().stream()
244 .<String>mapMulti((v, dump) -> {
245 switch (v) {
246 case Value val -> dump.accept(indexer.nameOf(val));
247 case Optional<?> o when o.isPresent() && o.get() instanceof Value val -> dump.accept(indexer.nameOf(val));
248 case List l -> l.forEach(val -> dump.accept(indexer.nameOf((Value)val)));
249 default -> dump.accept(""); // empty names for unused optional inputs
250 }
251 }).toList();
252 // trim trailing empty names
253 return inputNames.reversed().stream().dropWhile(String::isEmpty).toList().reversed();
254 }
255
256 static List<NodeProto> nodes(String domain, Indexer indexer, List<Op> ops) {
257 return ops.stream().<NodeProto>mapMulti((op, opNodes) -> {
258 switch (op) {
259 case OnnxOps.If ifOp ->
260 opNodes.accept(node(
261 ifOp.schema().name(),
262 List.of(indexer.nameOf(ifOp.operands().getFirst())),
263 IntStream.range(0, ifOp.resultType() instanceof TupleType tt ? tt.componentTypes().size() : 1).mapToObj(o -> indexer.nameOf(ifOp.result(), o)).toList(),
264 java.util.Map.of(
265 "then_branch", graph(domain, null, indexer, ifOp.thenBranch().entryBlock(), List.of(), 0),
266 "else_branch", graph(domain, null, indexer, ifOp.elseBranch().entryBlock(), List.of(), 0))));
267 case OnnxOps.Loop loopOp -> {
268 opNodes.accept(node(loopOp.schema().name(),
269 expandTuples(indexer, loopOp.operands()),
270 IntStream.range(0, loopOp.resultType() instanceof TupleType tt ? tt.componentTypes().size() : 1).mapToObj(o -> indexer.nameOf(loopOp.result(), o)).toList(),
271 java.util.Map.of(
272 "body", graph(domain, null, indexer, loopOp.loopBody().entryBlock(), List.of(), 2))));
273 }
274 case OnnxOp onnxOp ->
275 opNodes.accept(node(
276 onnxOp.schema().name(),
277 opInputNames(indexer, onnxOp.onnxInputs()),
278 IntStream.range(0, onnxOp.onnxOutputs().size()).mapToObj(o -> indexer.nameOf(onnxOp.result(), o)).toList(),
279 onnxOp.onnxAttributes()));
280 case CoreOp.FuncCallOp fco ->
281 opNodes.accept(node(
282 domain,
283 fco.funcName(),
284 expandTuples(indexer, fco.operands()),
285 expandTuples(indexer, List.of(fco.result())),
286 java.util.Map.of()));
287 case CoreOp.ReturnOp _, CoreOp.ConstantOp _ -> { // skip
288 }
289 case CoreOp.TupleLoadOp tlo ->
290 indexer.mapTupleLoad(tlo.result(), tlo.operands().getFirst(), tlo.index());
291 case CoreOp.TupleOp to ->
292 indexer.mapTupleElements(to.result(), to.operands());
293 case JavaOp.InvokeOp io when io.invokeReference().refType().equals(JavaType.type(List.class)) -> {
294 if (io.invokeReference().name().equals("get") && io.operands().getLast() instanceof Op.Result or && or.op() instanceof CoreOp.ConstantOp co && co.value() instanceof Integer i) {
295 indexer.mapTupleLoad(io.result(), io.operands().getFirst(), i);
296 } else if (io.invokeReference().name().equals("of")) {
297 indexer.mapTupleElements(io.result(), io.operands());
298 } else {
299 throw new UnsupportedOperationException(op.toText());
300 }
301 }
302 default -> {
303 throw new UnsupportedOperationException(op.toText());
304 }
305 }
306 }).toList();
307 }
308
309 static List<ValueInfoProto> tensorInfos(Indexer indexer, List<Block.Parameter> args, int scalarArgs) {
310 var infos = new ArrayList<ValueInfoProto>();
311 for (var arg : args) {
312 switch (arg.type()) {
313 case OnnxType.TensorType tt ->
314 infos.add(tensorInfo(indexer.nameOf(arg), tt.eType().id(), infos.size() < scalarArgs));
315 case TupleType tt -> {
316 var ct = tt.componentTypes();
317 for (int i = 0; i < ct.size(); i++) {
318 infos.add(tensorInfo(indexer.nameOf(arg, i), ((OnnxType.TensorType)ct.get(i)).eType().id(), infos.size() < scalarArgs));
319 }
320 }
321 default ->
322 throw new UnsupportedOperationException(arg.type().toString());
323 }
324 }
325 return infos;
326 }
327
328 static GraphProto graph(String name, List<TensorProto> initializers, List<ValueInfoProto> inputs, List<NodeProto> ops, List<String> outputNames) {
329 return new GraphProto()
330 .name(name)
331 .forEach(initializers, GraphProto::initializer)
332 .forEach(inputs, GraphProto::input)
333 .forEach(ops, GraphProto::node)
334 .forEach(outputNames, (g, oName) -> g.output(new ValueInfoProto().name(oName)));
335 }
336
337 static FunctionProto function(String functionDomain, List<String> imports, String functionName, List<String> inputNames, List<String> outputNames, List<NodeProto> ops) {
338 return new FunctionProto()
339 .domain(functionDomain)
340 .name(functionName)
341 .forEach(inputNames, FunctionProto::input)
342 .forEach(ops, FunctionProto::node)
343 .forEach(outputNames, FunctionProto::output)
344 .opsetImport(new OperatorSetIdProto().version(OPSET_VERSION))
345 .forEach(imports, (f, d) -> f.opsetImport(new OperatorSetIdProto().domain(d).version(1)));
346 }
347
348 static NodeProto node(String domain, String opName, List<String> inputNames, List<String> outputNames, java.util.Map<String, Object> attributes) {
349 return new NodeProto()
350 .domain(domain)
351 .opType(opName)
352 .forEach(inputNames, NodeProto::input)
353 .forEach(attributes.entrySet(), (n, ae) -> n.attribute(attribute(ae.getKey(), ae.getValue())))
354 .forEach(outputNames, NodeProto::output);
355 }
356
357 static NodeProto node(String opName, List<String> inputNames, List<String> outputNames, java.util.Map<String, Object> attributes) {
358 int di = opName.lastIndexOf('.');
359 return node(di < 0 ? null : opName.substring(0, di), opName.substring(di + 1), inputNames, outputNames, attributes);
360 }
361
362 static ValueInfoProto tensorInfo(String name, int tensorElementType) {
363 return tensorInfo(name, tensorElementType, false);
364 }
365
366 static ValueInfoProto tensorInfo(String name, int tensorElementType, boolean addScalarShape) {
367 var t = new TypeProto.Tensor().elemType(tensorElementType);
368 if (addScalarShape) t.shape(new TensorShapeProto());
369 return new ValueInfoProto()
370 .name(name)
371 .type(new TypeProto().tensorType(t));
372 }
373
374 static TensorProto tensorProto(String name, Tensor tensor, Function<Tensor, ExternalTensorDataInfo> tensorDataExternalizer) {
375 ExternalTensorDataInfo extInfo = tensorDataExternalizer.apply(tensor);
376 TensorProto tp = new TensorProto()
377 .name(name)
378 .dataType(tensor.elementType().id)
379 .dims(tensor.shape());
380 return extInfo == null
381 ? tp.rawData(tensor.data().toArray(ValueLayout.JAVA_BYTE))
382 : tp.externalData(new StringStringEntryProto().key("location").value(extInfo.location()))
383 .externalData(new StringStringEntryProto().key("offset").value(String.valueOf(extInfo.offset())))
384 .externalData(new StringStringEntryProto().key("length").value(String.valueOf(extInfo.length())))
385 .dataLocation(DataLocation.EXTERNAL);
386 }
387
388 static TensorProto tensorProto(Tensor tensor) {
389 return new TensorProto()
390 .dataType(tensor.elementType().id)
391 .dims(tensor.shape())
392 .rawData(tensor.data().toArray(ValueLayout.JAVA_BYTE));
393 }
394
395 static AttributeProto attribute(String name, Object value) {
396 var attr = new AttributeProto().name(name);
397 switch (value) {
398 case Float f -> {
399 attr.type(AttributeType.FLOAT).f(f);
400 }
401 case Long l -> {
402 attr.type(AttributeType.INT).i(l);
403 }
404 case GraphProto g -> {
405 attr.type(AttributeType.GRAPH).g(g.name(name));
406 }
407 case float[] floats -> {
408 attr.type(AttributeType.FLOATS);
409 attr.floats(floats);
410 }
411 case long[] longs -> {
412 attr.type(AttributeType.INTS);
413 attr.ints(longs);
414 }
415 case String s -> {
416 attr.type(AttributeType.STRING);
417 attr.s(s.getBytes());
418 }
419 case Tensor<?> t -> {
420 attr.type(AttributeType.TENSOR);
421 attr.t(tensorProto(t));
422 }
423 default -> {
424 throw new UnsupportedOperationException(value.getClass().toString()); // @@@ ToDo
425 }
426 }
427 return attr;
428 }
429 }