OMTensor.h

Go to the documentation of this file.

/*
 * SPDX-License-Identifier: Apache-2.0
 */
 
//===-------------- OMTensor.h - OMTensor Declaration header --------------===//
//
// Copyright 2019-2023 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 const int64_t *omTensorGetShape(const OMTensor *tensor);
 
OM_EXTERNAL_VISIBILITY void omTensorSetShape(
    OMTensor *tensor, const int64_t *shape);
 
OM_EXTERNAL_VISIBILITY const int64_t *omTensorGetStrides(
    const OMTensor *tensor);
 
OM_EXTERNAL_VISIBILITY void omTensorSetStrides(
    OMTensor *tensor, const int64_t *stride);
 
OM_EXTERNAL_VISIBILITY void omTensorSetStridesWithPyArrayStrides(
    OMTensor *tensor, const int64_t *stridesInBytes);
 
OM_EXTERNAL_VISIBILITY OM_DATA_TYPE omTensorGetDataType(const OMTensor *tensor);
 
OM_EXTERNAL_VISIBILITY void omTensorSetDataType(
    OMTensor *tensor, OM_DATA_TYPE dataType);
 
/* Helper function to get the ONNX data type size in bytes */
static inline int64_t getDataTypeSize(OM_DATA_TYPE dataType) {
  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