experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

switched from unsigned char to TF_Bool. Added input checking for TF_AllocatorAttributes

Browse Source
This commit is contained in:
Daniel Nguyen 2020-08-11 21:01:09 +00:00
parent 078052d2a9
commit a7e4df9246
6 changed files with 22 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/c/BUILD
View File

@ -81,6 +81,7 @@ tf_cuda_library(
hdrs = [
"c_api.h",
"c_api_internal.h",
"c_api_macros.h",
"tf_datatype.h",
"tf_tensor.h",
"tf_tstring.h",

6
tensorflow/c/c_api_macros.h
View File

@ -30,6 +30,12 @@ limitations under the License.
#endif // _WIN32
#endif // SWIG
// TF_Bool is the C API typedef for unsigned char, while TF_BOOL is
// the datatype for boolean tensors.
#ifndef TF_BOOL_DEFINED
#define TF_Bool unsigned char
#endif // TF_BOOL_DEFINED
// Macro used to calculate struct size for maintaining ABI stability across
// different struct implementations.
#ifndef TF_OFFSET_OF_END

14
tensorflow/c/kernels.cc
View File

@ -261,7 +261,6 @@ TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context, int index,
size_t len, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(context);
static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
"64-bit int types should match in size");
tensorflow::gtl::ArraySlice<tensorflow::int64> dimarray(
@ -286,22 +285,29 @@ TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context, TF_DataType dtype,
TF_AllocatorAttributes* attributes,
TF_Status* status) {
auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(context);
TF_SetStatus(status, TF_OK, "");
TF_SetStatus(status, TF_OK, "");
static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
"64-bit int types should match in size");
tensorflow::gtl::ArraySlice<tensorflow::int64> dimarray(
reinterpret_cast<tensorflow::int64*>(dims), num_dims);
if (attributes && !attributes->struct_size) {
TF_SetStatus(status, TF_INVALID_ARGUMENT, "TF_AllocatorAttributes struct "
"size member must be set to TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE");
return nullptr;
}
tensorflow::AllocatorAttributes allocator_attr;
if (attributes->on_host) {
if (attributes && attributes->on_host) {
allocator_attr.set_on_host(true);
}
tensorflow::Status s;
tensorflow::Tensor tensor;
TF_Tensor* tf_tensor;
s = cc_ctx->allocate_temp(static_cast<tensorflow::DataType>(dtype),
tensorflow::TensorShape(dimarray), &tensor, allocator_attr);
if (!s.ok()) {
::tensorflow::Set_TF_Status_from_Status(status, s);
return nullptr;
}
TF_Tensor* tf_tensor;
tf_tensor = TF_TensorFromTensor(tensor, &s);
if (!s.ok()) {
::tensorflow::Set_TF_Status_from_Status(status, s);

2
tensorflow/c/kernels.h
View File

@ -203,7 +203,7 @@ TF_CAPI_EXPORT TF_Tensor* TF_AllocateOutput(TF_OpKernelContext* context,
// Allocates a temporary Tensor of the specified type and shape. The
// Tensor must not be used after kernel construction is
// complete.
//
// num_dims must equal the size of array dims
TF_CAPI_EXPORT extern TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context,
TF_DataType dtype, int64_t* dims, int num_dims, TF_AllocatorAttributes*

7
tensorflow/c/kernels_test.cc
View File

@ -368,13 +368,12 @@ class DeviceKernelOpTest : public OpsTestBase {
#endif
};
// Helper function for tests that validates that the tensor has
// shape and type corresponding to dims and dtype.
// Validates that the tensor has shape and type corresponding to
// dims and dtype.
void validate_tensor(TF_Tensor* tensor, int64_t* dims, int64_t num_dims,
TF_DataType dtype);
// Helper function for tests that copies data of length
// tensor_size_bytes from values to tensor
// Copies data of length tensor_size_bytes from values to tensor.
template <typename T>
void set_tensor_data(TF_Tensor* tensor, T* values, size_t tensor_size_bytes,
TF_OpKernelContext* ctx);

2
tensorflow/c/tf_tensor.h
View File

@ -50,7 +50,7 @@ extern "C" {
typedef struct TF_AllocatorAttributes {
size_t struct_size;
// Set boolean to 1 for CPU allocation, else 0.
unsigned char on_host;
TF_Bool on_host;
} TF_AllocatorAttributes;