OMTensor.h

Go to the documentation of this file.

//===-------------- OMTensor.h - OMTensor Declaration header --------------===//
//
// Copyright 2019-2020 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

#ifdef __APPLE__
#include <stdlib.h>
#else
#include <malloc.h>
#endif // #ifdef __APPLE__

#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

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

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

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

void omTensorDestroy(OMTensor *tensor);

void *omTensorGetDataPtr(OMTensor *tensor);

int64_t *omTensorGetShape(OMTensor *tensor);

void omTensorSetShape(OMTensor *tensor, int64_t *shape);

int64_t *omTensorGetStrides(OMTensor *tensor);

void omTensorSetStrides(OMTensor *tensor, int64_t *stride);

void omTensorSetStridesWithPyArrayStrides(
    OMTensor *tensor, int64_t *stridesInBytes);

OM_DATA_TYPE omTensorGetDataType(OMTensor *tensor);

void omTensorSetDataType(OMTensor *tensor, OM_DATA_TYPE dataType);

/* Helper function to get the ONNX data type size in bytes */
static inline int getDataTypeSize(OM_DATA_TYPE dataType) {
  return OM_DATA_TYPE_SIZE[dataType];
}

int64_t omTensorGetBufferSize(OMTensor *tensor);

int omTensorGetRank(OMTensor *tensor);

int64_t omTensorGetNumElems(OMTensor *tensor);

int omTensorGetOwning(OMTensor *tensor);

#ifdef __cplusplus
}
#endif

#endif // ONNX_MLIR_OMTENSOR_H