From 54b76dc588d28643fbfd0a297870d79326b11abb Mon Sep 17 00:00:00 2001 From: Daniel Nguyen Date: Tue, 4 Aug 2020 22:08:35 +0000 Subject: [PATCH] switched from allocator attributes struct to opaque pointer --- tensorflow/c/kernels.cc | 6 +----- tensorflow/c/kernels.h | 1 - tensorflow/c/kernels_test.cc | 36 +++++++++++++------------------ tensorflow/c/tf_tensor.cc | 17 +++++++++++++++ tensorflow/c/tf_tensor.h | 22 +++++++++---------- tensorflow/c/tf_tensor_internal.h | 5 +++++ 6 files changed, 49 insertions(+), 38 deletions(-) diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc index 6d9d7d57517..096c8e41812 100644 --- a/tensorflow/c/kernels.cc +++ b/tensorflow/c/kernels.cc @@ -289,15 +289,11 @@ TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context, TF_DataType dtype, TF_SetStatus(status, TF_OK, ""); tensorflow::gtl::ArraySlice dimarray( reinterpret_cast(dims), num_dims); - tensorflow::AllocatorAttributes allocator_attr; - if (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(dtype), - tensorflow::TensorShape(dimarray), &tensor, allocator_attr); + tensorflow::TensorShape(dimarray), &tensor, attributes->alloc_attrs); if (!s.ok()) { ::tensorflow::Set_TF_Status_from_Status(status, s); return nullptr; diff --git a/tensorflow/c/kernels.h b/tensorflow/c/kernels.h index d6a7070de06..ee865613b6e 100644 --- a/tensorflow/c/kernels.h +++ b/tensorflow/c/kernels.h @@ -209,7 +209,6 @@ TF_CAPI_EXPORT extern TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context, TF_DataType dtype, int64_t* dims, int num_dims, TF_AllocatorAttributes* alloc_attrs, TF_Status* status); - #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/kernels_test.cc b/tensorflow/c/kernels_test.cc index 708c39690a8..5239274a2ce 100644 --- a/tensorflow/c/kernels_test.cc +++ b/tensorflow/c/kernels_test.cc @@ -467,16 +467,13 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempSizeOne) { // Allocate scalar TF_Tensor TF_Status* s = TF_NewStatus(); int64_t dim = 1; - TF_AllocatorAttributes alloc_attrs; - alloc_attrs.struct_size = TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE; -#if GOOGLE_CUDA - alloc_attrs.on_host = 0; -#else - alloc_attrs.on_host = 1; + TF_AllocatorAttributes* alloc_attrs = TF_NewAllocatorAttributes(); +#if !GOOGLE_CUDA + TF_AllocatorAttributesSetOnHost(alloc_attrs); #endif TF_Tensor* output = TF_AllocateTemp( /*context=*/ctx, /*dtype=*/TF_FLOAT, /*dims=*/&dim, - /*num_dims=*/1, /*allocator_attributes*/ &alloc_attrs, s); + /*num_dims=*/1, /*allocator_attributes*/ alloc_attrs, s); size_t tensor_size_bytes = TF_DataTypeSize(TF_FLOAT); EXPECT_EQ(TF_OK, TF_GetCode(s)); validate_tensor(output, &dim, 1, TF_FLOAT); @@ -487,6 +484,7 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempSizeOne) { TF_SetOutput(ctx, 0, output, s); TF_DeleteStatus(s); TF_DeleteTensor(output); + TF_DeleteAllocatorAttributes(alloc_attrs); }; SetupOp("AllocateTempOp1", "AllocateTemp1", my_compute_func); @@ -504,21 +502,19 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempEmpty) { TF_Status* s = TF_NewStatus(); // Allocate empty TF_Tensor int64_t dim = 0; - TF_AllocatorAttributes alloc_attrs; - alloc_attrs.struct_size = TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE; -#if GOOGLE_CUDA - alloc_attrs.on_host = 0; -#else - alloc_attrs.on_host = 1; + TF_AllocatorAttributes* alloc_attrs = TF_NewAllocatorAttributes(); +#if !GOOGLE_CUDA + TF_AllocatorAttributesSetOnHost(alloc_attrs); #endif TF_Tensor* output = TF_AllocateTemp( /*context=*/ctx, /*dtype=*/TF_FLOAT, /*dims=*/&dim, - /*num_dims=*/1, /*allocator_attributes*/ &alloc_attrs, s); + /*num_dims=*/1, /*allocator_attributes*/ alloc_attrs, s); EXPECT_EQ(TF_OK, TF_GetCode(s)); validate_tensor(output, &dim, 1, TF_FLOAT); TF_SetOutput(ctx, 0, output, s); TF_DeleteStatus(s); TF_DeleteTensor(output); + TF_DeleteAllocatorAttributes(alloc_attrs); }; SetupOp("AllocateTempOp0", "AllocateTemp0", my_compute_func); @@ -537,16 +533,13 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempSize2x3) { size_t tensor_size_bytes = 6 * TF_DataTypeSize(TF_FLOAT); // Allocate 2x3 TF_Tensor int64_t dim[2] = {2, 3}; - TF_AllocatorAttributes alloc_attrs; - alloc_attrs.struct_size = TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE; -#if GOOGLE_CUDA - alloc_attrs.on_host = 0; -#else - alloc_attrs.on_host = 1; + TF_AllocatorAttributes* alloc_attrs = TF_NewAllocatorAttributes(); +#if !GOOGLE_CUDA + TF_AllocatorAttributesSetOnHost(alloc_attrs); #endif TF_Tensor* output = TF_AllocateTemp( /*context=*/ctx, /*dtype=*/TF_FLOAT, /*dims=*/dim, - /*num_dims=*/2, /*allocator_attributes*/ &alloc_attrs, s); + /*num_dims=*/2, /*allocator_attributes*/ alloc_attrs, s); EXPECT_EQ(TF_OK, TF_GetCode(s)); validate_tensor(output, dim, 2, TF_FLOAT); @@ -556,6 +549,7 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempSize2x3) { TF_SetOutput(ctx, 0, output, s); TF_DeleteStatus(s); TF_DeleteTensor(output); + TF_DeleteAllocatorAttributes(alloc_attrs); }; SetupOp("AllocateTempOp2x3", "AllocateTempOp2x3", my_compute_func); diff --git a/tensorflow/c/tf_tensor.cc b/tensorflow/c/tf_tensor.cc index 0feb986ce44..12405939bfc 100644 --- a/tensorflow/c/tf_tensor.cc +++ b/tensorflow/c/tf_tensor.cc @@ -321,3 +321,20 @@ bool TensorInterface::IsAligned() const { return tensor_.IsAligned(); } } // namespace tensorflow bool TF_TensorIsAligned(const TF_Tensor* t) { return t->tensor->IsAligned(); } + +TF_AllocatorAttributes* TF_NewAllocatorAttributes() { + return new TF_AllocatorAttributes{tensorflow::AllocatorAttributes()}; +} + +void TF_AllocatorAttributesSetOnHost(TF_AllocatorAttributes* tf_alloc_attrs) { + tf_alloc_attrs->alloc_attrs.set_on_host(true); +} + +void TF_DeleteAllocatorAttributes(TF_AllocatorAttributes* tf_alloc_attrs) { + if (tf_alloc_attrs == nullptr) { + return; + } + else { + delete tf_alloc_attrs; + } +} diff --git a/tensorflow/c/tf_tensor.h b/tensorflow/c/tf_tensor.h index ced57df77d4..9a043ce0538 100644 --- a/tensorflow/c/tf_tensor.h +++ b/tensorflow/c/tf_tensor.h @@ -46,17 +46,6 @@ limitations under the License. extern "C" { #endif -// Allocator Attributes used for tensor allocation. -typedef struct TF_AllocatorAttributes { - size_t struct_size; - // Set boolean to 1 for CPU allocation, else 0. - unsigned char on_host; -} TF_AllocatorAttributes; - - -#define TF_ALLOCATOR_ATTRIBUTES_STRUCT_SIZE \ - TF_OFFSET_OF_END(TF_AllocatorAttributes, on_host) - // -------------------------------------------------------------------------- // TF_Tensor holds a multi-dimensional array of elements of a single data type. // For all types other than TF_STRING, the data buffer stores elements @@ -163,6 +152,17 @@ TF_CAPI_EXPORT extern void TF_TensorBitcastFrom(const TF_Tensor* from, // Returns bool iff this tensor is aligned. TF_CAPI_EXPORT extern bool TF_TensorIsAligned(const TF_Tensor*); +// Allocator Attributes used for tensor allocation. +typedef struct TF_AllocatorAttributes TF_AllocatorAttributes; + +TF_CAPI_EXPORT extern TF_AllocatorAttributes* TF_NewAllocatorAttributes(); + +TF_CAPI_EXPORT extern void TF_AllocatorAttributesSetOnHost( + TF_AllocatorAttributes* tf_alloc_attrs); + +TF_CAPI_EXPORT extern void TF_DeleteAllocatorAttributes( + TF_AllocatorAttributes* tf_alloc_attrs); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/tf_tensor_internal.h b/tensorflow/c/tf_tensor_internal.h index 7a896dc5d11..71cac0afeb1 100644 --- a/tensorflow/c/tf_tensor_internal.h +++ b/tensorflow/c/tf_tensor_internal.h @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/core/framework/allocation_description.pb.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/platform/casts.h" // Internal structures used by the C API. These are likely to change and should @@ -124,4 +125,8 @@ Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst); TF_Tensor* TF_TensorFromTensor(const Tensor& src, Status* status); } // namespace tensorflow +typedef struct TF_AllocatorAttributes { + tensorflow::AllocatorAttributes alloc_attrs; +} TF_AllocatorAttributes; + #endif // TENSORFLOW_C_TF_TENSOR_INTERNAL_H_