diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc index 096c8e41812..6d9d7d57517 100644 --- a/tensorflow/c/kernels.cc +++ b/tensorflow/c/kernels.cc @@ -289,11 +289,15 @@ 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, attributes->alloc_attrs); + tensorflow::TensorShape(dimarray), &tensor, allocator_attr); 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 ee865613b6e..d6a7070de06 100644 --- a/tensorflow/c/kernels.h +++ b/tensorflow/c/kernels.h @@ -209,6 +209,7 @@ 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 5239274a2ce..708c39690a8 100644 --- a/tensorflow/c/kernels_test.cc +++ b/tensorflow/c/kernels_test.cc @@ -467,13 +467,16 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempSizeOne) { // Allocate scalar TF_Tensor TF_Status* s = TF_NewStatus(); int64_t dim = 1; - TF_AllocatorAttributes* alloc_attrs = TF_NewAllocatorAttributes(); -#if !GOOGLE_CUDA - TF_AllocatorAttributesSetOnHost(alloc_attrs); + 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; #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); @@ -484,7 +487,6 @@ 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); @@ -502,19 +504,21 @@ TEST_F(DeviceKernelOpTest, TestAllocateTempEmpty) { TF_Status* s = TF_NewStatus(); // Allocate empty TF_Tensor int64_t dim = 0; - TF_AllocatorAttributes* alloc_attrs = TF_NewAllocatorAttributes(); -#if !GOOGLE_CUDA - TF_AllocatorAttributesSetOnHost(alloc_attrs); + 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; #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); @@ -533,13 +537,16 @@ 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 = TF_NewAllocatorAttributes(); -#if !GOOGLE_CUDA - TF_AllocatorAttributesSetOnHost(alloc_attrs); + 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; #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); @@ -549,7 +556,6 @@ 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 12405939bfc..0feb986ce44 100644 --- a/tensorflow/c/tf_tensor.cc +++ b/tensorflow/c/tf_tensor.cc @@ -321,20 +321,3 @@ 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 9a043ce0538..ced57df77d4 100644 --- a/tensorflow/c/tf_tensor.h +++ b/tensorflow/c/tf_tensor.h @@ -46,6 +46,17 @@ 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 @@ -152,17 +163,6 @@ 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 71cac0afeb1..7a896dc5d11 100644 --- a/tensorflow/c/tf_tensor_internal.h +++ b/tensorflow/c/tf_tensor_internal.h @@ -23,7 +23,6 @@ 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 @@ -125,8 +124,4 @@ 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_