Implementation:Microsoft Onnxruntime OnnxSparseTensor JNI

Knowledge Sources	Description
Source File	java/src/main/native/ai_onnxruntime_OnnxSparseTensor.c
Repository	Microsoft/onnxruntime

Domains

• JNI Native Bridge
• Sparse Tensor Operations
• ORT Value Management

Overview

ai_onnxruntime_OnnxSparseTensor.c is the JNI native implementation file for the OnnxSparseTensor Java class. It provides native methods for creating sparse tensors in COO, CSRC (CSR), and Block Sparse formats from Java direct buffers, extracting index and value buffers from existing sparse tensors, querying shapes, and releasing native OrtValue resources.

Description

Key function groups:

• Buffer extraction: getIndicesBuffer, getInnerIndicesBuffer, and getValuesBuffer return direct ByteBuffers wrapping the native memory of sparse tensor components. The indices format is determined by the sparse format (COO, CSRC, Block Sparse).
• Shape queries: getIndicesShape, getInnerIndicesShape, and getValuesShape extract dimension arrays from the sparse tensor type/shape info, returning long[] to Java.
• Tensor creation: createCSRCSparseTensorFromBuffer creates a CSR/CSRC sparse tensor from three direct buffers (outer indices, inner indices, and values) plus dense and values shapes. createSparseTensorFromBuffer creates COO or Block Sparse tensors from indices and data buffers, dispatching to UseCooIndices or UseBlockSparseIndices based on the sparsity type.
• Resource cleanup: close releases the native OrtValue via api->ReleaseValue.

All functions use checkOrtStatus for error handling and throw OrtException on failure. The inner indices buffer is only available for CSRC format; requesting it for other formats throws an exception.

Code Reference

Source Location

// File: java/src/main/native/ai_onnxruntime_OnnxSparseTensor.c

Signature

// Buffer extraction
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getIndicesBuffer
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getInnerIndicesBuffer
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getValuesBuffer
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);

// Shape queries
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getIndicesShape
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getInnerIndicesShape
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);
JNIEXPORT jobject JNICALL Java_ai_onnxruntime_OnnxSparseTensor_getValuesShape
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);

// Tensor creation
JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OnnxSparseTensor_createCSRCSparseTensorFromBuffer
    (JNIEnv*, jclass, jlong apiHandle, jlong allocatorHandle,
     jobject indicesBuffer, jint indicesBufferPos, jlong indicesBufferSize,
     jobject innerIndicesBuffer, jint innerIndicesBufferPos, jlong innerIndicesBufferSize,
     jobject dataBuffer, jint dataBufferPos,
     jlongArray denseShape, jlongArray valuesShape, jint onnxTypeJava);
JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OnnxSparseTensor_createSparseTensorFromBuffer
    (JNIEnv*, jclass, jlong apiHandle, jlong allocatorHandle,
     jobject indicesBuffer, jint indicesBufferPos, jlong indicesBufferSize,
     jobject dataBuffer, jint dataBufferPos,
     jlongArray denseShape, jlongArray indicesShape, jlongArray valuesShape,
     jint onnxTypeJava, jint sparsityTypeJava);

// Resource cleanup
JNIEXPORT void JNICALL Java_ai_onnxruntime_OnnxSparseTensor_close
    (JNIEnv*, jobject, jlong apiHandle, jlong handle);

Import

#include "OrtJniUtil.h"
#include "ai_onnxruntime_OnnxSparseTensor.h"

I/O Contract

Inputs

Name	Type	Description
apiHandle	jlong	Pointer to the ORT C API struct
handle	jlong	Pointer to the native OrtValue (sparse tensor)
allocatorHandle	jlong	Pointer to the ORT allocator
indicesBuffer	jobject (Buffer)	Direct buffer containing sparse indices data
dataBuffer	jobject (Buffer)	Direct buffer containing sparse tensor values
denseShape	jlongArray	Shape of the dense representation
valuesShape	jlongArray	Shape of the values array
onnxTypeJava	jint	ONNX tensor element data type as Java int
sparsityTypeJava	jint	Sparse format type (COO, Block Sparse) as Java int

Outputs

Name	Type	Description
ByteBuffer	jobject	Direct ByteBuffer wrapping native index/value memory
long[]	jlongArray	Shape dimensions of the indices or values
nativeHandle	jlong	Pointer to the newly created OrtValue (0 on failure)

Usage Examples

// Inside JNI: get the indices buffer from a COO sparse tensor
jobject indicesBuf = Java_ai_onnxruntime_OnnxSparseTensor_getIndicesBuffer(
    jniEnv, sparseObj, apiHandle, ortValueHandle);

// Inside JNI: create a CSRC sparse tensor from buffers
jlong newHandle = Java_ai_onnxruntime_OnnxSparseTensor_createCSRCSparseTensorFromBuffer(
    jniEnv, cls, apiHandle, allocHandle,
    outerIndicesBuf, 0, outerSize,
    innerIndicesBuf, 0, innerSize,
    dataBuf, 0,
    denseShapeArr, valuesShapeArr, onnxType);

// Release the sparse tensor when done
Java_ai_onnxruntime_OnnxSparseTensor_close(jniEnv, sparseObj, apiHandle, ortValueHandle);

Related Pages

• OnnxSparseTensor.java - Java-side sparse tensor class calling these natives
• OrtJniUtil.c - Shared utility functions (checkOrtStatus, convertToONNXDataFormat, onnxTypeSize)
• ai_onnxruntime_OnnxTensor.c - Dense tensor JNI counterpart