Revert "switched from allocator attributes struct to opaque pointer"

This reverts commit 54b76dc588.
This commit is contained in:
Daniel Nguyen 2020-08-05 19:13:52 +00:00
parent 54b76dc588
commit b6960d8bd1
6 changed files with 38 additions and 49 deletions

View File

@ -289,11 +289,15 @@ TF_Tensor* TF_AllocateTemp(TF_OpKernelContext* context, TF_DataType dtype,
TF_SetStatus(status, TF_OK, "");
tensorflow::gtl::ArraySlice<tensorflow::int64> dimarray(
reinterpret_cast<tensorflow::int64*>(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<tensorflow::DataType>(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;

View File

@ -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

View File

@ -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);

View File

@ -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;
}
}

View File

@ -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

View File

@ -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_