com.ibm.onnxmlir.OMModel Class Reference

Static Public Member Functions
static OMTensorList	mainGraph (OMTensorList list)
static String[]	queryEntryPoints ()
static String	inputSignature ()
static String	inputSignature (String entry_point)
static String	outputSignature ()
static String	outputSignature (String entry_point)

Detailed Description

ONNX-MLIR Runtime Java API documentation

Introduction

ONNX-MLIR project comes with an executable onnx-mlir capable of compiling onnx models to a jar. In this documentation, we demonstrate how to interact programmatically with the compiled jar using ONNX-MLIR's Java Runtime API.

Java Runtime API

Classes

OMModel is the class implementing the default model entry point and input/output signature functions.

OMTensor is the class used to describe the runtime information (rank, shape, data type, etc) associated with a tensor input or output.

OMTensorList is the class used to hold an array of OMTensor so that they can be passed into and out of the compiled model as inputs and outputs.

Model Entry Point Signature

All compiled models will have the same exact Java function signature equivalent to:

public static OMTensorList mainGraph(OMTensorList list)

Intuitively, the model takes a list of tensors as input and returns a list of tensors as output.

Invoke Models Using Java Runtime API

We demonstrate using the API functions to run a simple ONNX model consisting of an add operation. To create such an onnx model, use this python script

To compile the above model, run onnx-mlir --EmitJNI add.onnx and a jar file "add.jar" should appear. We can use the following Java code to call into the compiled function computing the sum of two inputs:

import com.ibm.onnxmlir.OMModel;
import com.ibm.onnxmlir.OMTensor;
import com.ibm.onnxmlir.OMTensorList;
 
public class Add {
  public static void main(String[] args) {
    // Shared shape.
    long[] shape = {3, 2};
    // Construct x1 omt filled with 1.
    float[] x1Data = {1.f, 1.f, 1.f, 1.f, 1.f, 1.f};
    OMTensor x1 = new OMTensor(x1Data, shape);
    // Construct x2 omt filled with 2.
    float[] x2Data = {2.f, 2.f, 2.f, 2.f, 2.f, 2.f};
    OMTensor x2 = new OMTensor(x2Data, shape);
    // Construct a list of omts as input.
    OMTensor[] list = new OMTensor[]{x1, x2};
    OMTensorList input = new OMTensorList(list);
    // Call the compiled onnx model function.
    OMTensorList output = OMModel.mainGraph(input);
    // Get the first omt as output.
    OMTensor[] y = output.getOmtArray();
    float[] outputPtr = y[0].getFloatData();
    // Print its content, should be all 3.
    for (int i = 0; i < 6; i++)
      System.out.print(outputPtr[i] + " ");
  }
}

Compile with javac -cp javaruntime.jar Add.java, you should see a class Add.class appearing. Run it with java -cp .:add.jar Add, and the output should be:

3.0 3.0 3.0 3.0 3.0 3.0

Exactly as it should be.

Member Function Documentation

inputSignature() [1/2]

static String com.ibm.onnxmlir.OMModel.inputSignature ( )

inlinestatic

Input signature of default model runtime entry point

Returns

JSON string of input signature

inputSignature() [2/2]

static String com.ibm.onnxmlir.OMModel.inputSignature ( String  entry_point )

inlinestatic

mainGraph()

static OMTensorList com.ibm.onnxmlir.OMModel.mainGraph ( OMTensorList  list )

inlinestatic

Default model runtime entry point

Parameters

list	input tensor list

Returns

output tensor list

outputSignature() [1/2]

static String com.ibm.onnxmlir.OMModel.outputSignature ( )

inlinestatic

Output signature of default model runtime entry point

Returns

JSON string of output signature

outputSignature() [2/2]

static String com.ibm.onnxmlir.OMModel.outputSignature ( String  entry_point )

inlinestatic

queryEntryPoints()

static String [] com.ibm.onnxmlir.OMModel.queryEntryPoints ( )

inlinestatic

Query all entry point names in the model.

Returns

String array of entry point names

---

The documentation for this class was generated from the following file:

• onnx-mlir/src/Runtime/jni/src/com/ibm/onnxmlir/OMModel.java