OMTensor.h

Go to the documentation of this file.

/*
 * SPDX-License-Identifier: Apache-2.0
 */
 
//===-------------- OMTensor.h - OMTensor Declaration header --------------===//
//
// Copyright 2019-2026 The IBM Research Authors.
//
// =============================================================================
//
// This file contains declaration of OMTensor data structures and
// API functions.
//
//===----------------------------------------------------------------------===//
 
#ifndef ONNX_MLIR_OMTENSOR_H
#define ONNX_MLIR_OMTENSOR_H
 
#ifdef __cplusplus
#include <algorithm>
#include <iostream>
#include <map>
#include <numeric>
#include <string>
#include <vector>
#else
#include <stdbool.h>
#endif // #ifdef __cplusplus
 
#if defined(__APPLE__) || defined(__MVS__)
#include <stdlib.h>
#else
#include <malloc.h>
#endif // #ifdef __APPLE__
 
#include "onnx-mlir/Compiler/OMCompilerMacros.h"
#include "onnx-mlir/Runtime/OnnxDataType.h"
 
/* Typically, MemRefs in MLIR context are used as a compile-time constructs.
 * Information such as element type and rank of the data payload is statically
 * encoded, meaning that they are determined and fixed at compile-time. This
 * presents significant burden for any runtime components trying to interact
 * with the compiled executable.
 *
 * Thus a version of MemRef struct that is amenable to runtime manipulation is
 * provided as a basis for building any runtime-related components providing
 * user-facing programming interfaces. All information are dynamically encoded
 * as members of this struct so that they can be accessed and modified easily
 * during runtime.
 *
 * We will refer to it as an OMTensor.
 */
 
struct OMTensor;
 
#ifndef __cplusplus
typedef struct OMTensor OMTensor;
#endif
 
#ifdef __cplusplus
extern "C" {
#endif

OM_EXTERNAL_VISIBILITY OMTensor *omTensorCreate(
    void *data_ptr, const int64_t *shape, int64_t rank, OM_DATA_TYPE dtype);

OM_EXTERNAL_VISIBILITY OMTensor *omTensorCreateWithOwnership(void *data_ptr,
    const int64_t *shape, int64_t rank, OM_DATA_TYPE dtype, int64_t owning);

OM_EXTERNAL_VISIBILITY OMTensor *omTensorCreateEmpty(
    const int64_t *shape, int64_t rank, OM_DATA_TYPE dtype);

OM_EXTERNAL_VISIBILITY void omTensorDestroy(OMTensor *tensor);

OM_EXTERNAL_VISIBILITY void *omTensorGetDataPtr(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY void *omTensorGetAllocatedPtr(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY const int64_t *omTensorGetShape(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY OM_DATA_TYPE omTensorGetDataType(const OMTensor *tensor);
 
/* Helper function to get the ONNX data type size in bytes */
 
/* Note that this function does really support sub-byte quantities
 * like int4, unit4 and smaller, as we are returning here byte size
 * quantities. Eventually, we need to transform this (and the associated
 * OM_DATA_TYPE_SIZE) to return bit sizes so as to better represent such data
 * sizes. It also forces all functions in this library to either up-convert
 * sub-byte size to byte size values to process them.
 *
 */
static inline int64_t getDataTypeSize(OM_DATA_TYPE dataType) {
  /* Prevent out of bound access to OM_DATA_TYPE_SIZE */
  if ((int)dataType < 0 || (int)dataType > ONNX_TYPE_LAST)
    return OM_DATA_TYPE_SIZE[ONNX_TYPE_UNDEFINED];
  return OM_DATA_TYPE_SIZE[dataType];
}

OM_EXTERNAL_VISIBILITY int64_t omTensorGetBufferSize(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY int64_t omTensorGetRank(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY int64_t omTensorGetNumElems(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY int64_t omTensorGetOwning(const OMTensor *tensor);

OM_EXTERNAL_VISIBILITY void omTensorSetOwning(OMTensor *tensor, int64_t owning);

OM_EXTERNAL_VISIBILITY void omTensorPrint(
    const char *msg, const OMTensor *tensor);
 
#ifdef __cplusplus
}
#endif
 
#endif // ONNX_MLIR_OMTENSOR_H