wehub-resource-sync
193 lines
9.4 KiB
C
193 lines
9.4 KiB
C
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#ifndef TENSORFLOW_C_KERNELS_EXPERIMENTAL_H_
|
|
#define TENSORFLOW_C_KERNELS_EXPERIMENTAL_H_
|
|
|
|
#include "tensorflow/c/c_api_macros.h"
|
|
#include "tensorflow/c/kernels.h"
|
|
|
|
// --------------------------------------------------------------------------
|
|
// Experimental kernel C API for TensorFlow.
|
|
//
|
|
// The API here is subject to changes in the future.
|
|
// --------------------------------------------------------------------------
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
typedef struct TF_VariableInputLockHolder TF_VariableInputLockHolder;
|
|
|
|
// Expose higher level Assignment operation for Pluggable vendors to implement
|
|
// in the plugin for Training. The API takes in the context with indices for
|
|
// the input and value tensors. It also accepts the copy callback provided by
|
|
// pluggable vendor to do the copying of the tensors. The caller takes ownership
|
|
// of the `source` and `dest` tensors and is responsible for freeing them with
|
|
// TF_DeleteTensor. This function will return an error when the following
|
|
// conditions are met:
|
|
// 1. `validate_shape` is set to `true`
|
|
// 2. The variable is initialized
|
|
// 3. The shape of the value tensor doesn't match the shape of the variable
|
|
// tensor.
|
|
TF_CAPI_EXPORT extern void TF_AssignVariable(
|
|
TF_OpKernelContext* ctx, int input_index, int value_index,
|
|
bool validate_shape,
|
|
void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
|
|
TF_Tensor* dest),
|
|
TF_Status* status);
|
|
|
|
// Expose higher level Assignment operation for Pluggable vendors to implement
|
|
// in the plugin for Training on ref variables. The API takes in the context
|
|
// with indices for the input and value tensors. It also accepts the copy
|
|
// callback provided by pluggable vendor to do the copying of the tensors. The
|
|
// caller takes ownership of the `source` and `dest` tensors and is responsible
|
|
// for freeing them with TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_AssignRefVariable(
|
|
TF_OpKernelContext* ctx, int input_ref_index, int output_ref_index,
|
|
int value_index, bool use_locking, bool validate_shape,
|
|
void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
|
|
TF_Tensor* dest),
|
|
TF_Status* status);
|
|
|
|
// Expose higher level AssignUpdate operation for Pluggable vendors to implement
|
|
// in the plugin for Training. The API takes in the context with indices for the
|
|
// input and value tensors. It also accepts the copy callback provided by
|
|
// pluggable vendor to do the copying of the tensors and the update callback to
|
|
// apply the arithmetic operation. The caller takes ownership of the `source`,
|
|
// `dest`, `tensor` and `value` tensors and is responsible for freeing them with
|
|
// TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_AssignUpdateVariable(
|
|
TF_OpKernelContext* ctx, int input_index, int value_index, int Op,
|
|
int isVariantType,
|
|
void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
|
|
TF_Tensor* dest),
|
|
void (*updateFunc)(TF_OpKernelContext* ctx, TF_Tensor* tensor,
|
|
TF_Tensor* value, int Op),
|
|
TF_Status* status);
|
|
|
|
// Expose higher level temporary variable operator for Pluggable vendors to
|
|
// implement in the plugin for managing temporary variables. The API takes in
|
|
// the context with indices for the input and value tensors. It also accepts the
|
|
// allocator provided by pluggable vendor to do the allocate_temp of the
|
|
// tensors. The caller takes ownership of temporary variables and is responsible
|
|
// for freeing them with TF_DestroyTemporaryVariable. This function will return
|
|
// an error when the following conditions are met:
|
|
// 1. Cannot allocate a new temporary variable
|
|
// 2. Calling plugin allocator failed
|
|
TF_CAPI_EXPORT extern void TF_TemporaryVariable(
|
|
TF_OpKernelContext* ctx, TF_DataType dtype, const int64_t* dims,
|
|
int num_dims, TF_StringView* var_name,
|
|
void (*plugin_allocator)(TF_OpKernelContext*, TF_Tensor*, TF_DataType,
|
|
const int64_t*, int, TF_Status*),
|
|
TF_Status* tf_status);
|
|
|
|
// Expose higher level temporary variable operator for Pluggable vendors to
|
|
// implement in the plugin for destroying temporary variables. The API takes in
|
|
// the context with indices for the input and variable name. This function will
|
|
// return an error when either of the following conditions is met:
|
|
// 1. `input data type` is not ref type
|
|
// 2. Cannot find temporary variable by name in arguments
|
|
TF_CAPI_EXPORT extern void TF_DestroyTemporaryVariable(TF_OpKernelContext* ctx,
|
|
const int index,
|
|
TF_StringView* var_name,
|
|
TF_Status* tf_status);
|
|
|
|
// This is a helper function which acquires mutexes in-order to provide
|
|
// thread-safe way of performing weights update during the optimizer op. It
|
|
// returns an opaque LockHolder handle back to plugin. This handle is passed to
|
|
// the Release API for releasing the locks when the weight update is done. The
|
|
// caller takes ownership of the `source` and `dest` tensors and is responsible
|
|
// for freeing them with TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_MaybeLockVariableInputMutexesInOrder(
|
|
TF_OpKernelContext* ctx, bool do_lock, bool sparse, const int* const inputs,
|
|
size_t len,
|
|
void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
|
|
TF_Tensor* dest),
|
|
TF_VariableInputLockHolder** lockHolder, TF_Status* status);
|
|
|
|
// This interface returns `out` tensor which is updated corresponding to the
|
|
// variable passed with input index. The caller takes ownership of the `source`
|
|
// and `dest` tensors and is responsible for freeing them with TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_GetInputTensorFromVariable(
|
|
TF_OpKernelContext* ctx, int input, bool lock_held, bool isVariantType,
|
|
bool sparse,
|
|
void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,
|
|
TF_Tensor* dest),
|
|
TF_Tensor** out, TF_Status* status);
|
|
|
|
// This interface forwards the reference from input to the output tensors
|
|
// corresponding to the indices provided with `input_index` and `output_index`
|
|
TF_CAPI_EXPORT extern void TF_OpKernelContext_ForwardRefInputToRefOutput(
|
|
TF_OpKernelContext* ctx, int32_t input_index, int32_t output_index);
|
|
|
|
// The API releases the opaque lock handle returned with
|
|
// `TF_MaybeLockVariableInputMutexesInOrder` API
|
|
TF_CAPI_EXPORT extern void TF_ReleaseVariableInputLockHolder(
|
|
TF_VariableInputLockHolder* lockHolder);
|
|
|
|
// Allows plugin to get TF_Tensor when passed its input_name
|
|
TF_CAPI_EXPORT extern void TF_GetInputByName(TF_OpKernelContext* ctx,
|
|
const char* inputName,
|
|
TF_Tensor** tensor,
|
|
TF_Status* status);
|
|
|
|
// Interprets the named kernel construction attribute as a shape attribute and
|
|
// fills in `vals` with the size of each dimension. `vals` must point to an
|
|
// array of length at least `max_values` (ideally set to total_size from
|
|
// TF_OpKernelConstruction_GetAttrSize(ctx, attr_name, &list_size,
|
|
// &total_size)).
|
|
TF_CAPI_EXPORT extern void TF_OpKernelConstruction_GetAttrTensorShape(
|
|
TF_OpKernelConstruction* ctx, const char* attr_name, int64_t* dims,
|
|
size_t num_dims, TF_Status* status);
|
|
|
|
TF_CAPI_EXPORT extern bool TF_IsRefInput(TF_OpKernelContext* ctx, int i,
|
|
TF_Status* status);
|
|
|
|
#ifndef IS_MOBILE_PLATFORM
|
|
// Expose higher level AddN operation for Pluggable vendors to implement
|
|
// in the plugin for Variant data types. The API takes in the context and a
|
|
// callback provided by pluggable vendor to do a Binary Add operation on the
|
|
// tensors unwrapped from the Variant tensors. The caller takes ownership of the
|
|
// `a`, `b` and `out` tensors and is responsible for freeing them with
|
|
// TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_AddNVariant(
|
|
TF_OpKernelContext* ctx,
|
|
void (*binary_add_func)(TF_OpKernelContext* ctx, TF_Tensor* a, TF_Tensor* b,
|
|
TF_Tensor* out),
|
|
TF_Status* status);
|
|
|
|
// Expose higher level ZerosLike operation for Pluggable vendors to implement
|
|
// in the plugin for Variant data types. The API takes in the context and a
|
|
// callback provided by pluggable vendor to do a ZerosLike operation on the
|
|
// tensors unwrapped from the Variant tensors. The caller takes ownership of the
|
|
// `input` and `out` tensors and is responsible for freeing them with
|
|
// TF_DeleteTensor.
|
|
TF_CAPI_EXPORT extern void TF_ZerosLikeVariant(
|
|
TF_OpKernelContext* ctx,
|
|
void (*zeros_like_func)(TF_OpKernelContext* ctx, TF_Tensor* input,
|
|
TF_Tensor* out),
|
|
TF_Status* status);
|
|
|
|
typedef struct TF_CoordinationServiceAgent TF_CoordinationServiceAgent;
|
|
|
|
#endif // IS_MOBILE_PLATFORM
|
|
|
|
#ifdef __cplusplus
|
|
} /* end extern "C" */
|
|
#endif
|
|
|
|
#endif // TENSORFLOW_C_KERNELS_EXPERIMENTAL_H_
|