[code-reflection] Integrated: Very basic ONNX transformer

Paul Sandoz psandoz at openjdk.org
Wed Feb 5 23:40:39 UTC 2025


On Wed, 5 Feb 2025 23:18:14 GMT, Paul Sandoz <psandoz at openjdk.org> wrote:

> An ONNX transformer that transforms the Java code model of a Java ONNX function to an ONNX code model, from which we can create an ONNX model. This is just the first attempt and it will likely take many iterations to get this right.
> 
> There are three stages to the transformation:
> 1. Partially evaluate the Java code model, using a simple interpreter. Any operation that depends on a function input parameter, directly or indirectly, is ignored. This leaves operations that depend on constants, directly or indirectly, and since attributes are constants we can evaluate them and associate a list of attributes with each invocation to an ONNX operator.
> 2. Transform the Java code model to an ONNX model, only transforming operations whose results have not been evaluated namely all operations that depend, directly or indirectly on function inputs. An invocation to an ONNX operator is transformed to an instance of the corresponding ONNX operation, and corresponding attributes evaluated in the first step are passed directly as arguments to the operation factory method.
> 3. Clean up the ONNX model, transforming to pure SSA, and removing any operation whose result is unused.
> 
> So far we avoided any need to compute tensor shape information and we may not need it, although it would be nice to have if we could defer to the ONNX runtime to compute for each operation.
> 
> Currently we assume that the Java ONNX function accepts weights as parameters. This simplifies the transformation and modeling. When we export the ONNX code model a prefix of tensors can be given for the weights. We can revisit this modeling later on and improve the developer experience.
> 
> This PR also adjusts the mapping of ONNX attribute type to Java class. `INT` should map to `long`, and `INTS` to `long[]`.

This pull request has now been integrated.

Changeset: 408e7b56
Author:    Paul Sandoz <psandoz at openjdk.org>
URL:       https://git.openjdk.org/babylon/commit/408e7b56db440349a7d11871a84a91a372b75f8f
Stats:     2090 lines in 8 files changed: 897 ins; 22 del; 1171 mod

Very basic ONNX transformer

-------------

PR: https://git.openjdk.org/babylon/pull/308


More information about the babylon-dev mailing list