-
Notifications
You must be signed in to change notification settings - Fork 74k
/
kernels_experimental.h
192 lines (166 loc) · 9.43 KB
/
kernels_experimental.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
/* 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_