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

switched from allocator attributes struct to opaque pointer

Browse Source
This commit is contained in:
Daniel Nguyen 2020-08-04 22:08:35 +00:00
parent 8ef70baf37
commit 54b76dc588
6 changed files with 49 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/c/kernels.cc
View File

@ -289,15 +289,11 @@ 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, allocator_attr);
tensorflow::TensorShape(dimarray), &tensor, attributes->alloc_attrs);
if (!s.ok()) {
::tensorflow::Set_TF_Status_from_Status(status, s);
return nullptr;

1
tensorflow/c/kernels.h
View File

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

36
tensorflow/c/kernels_test.cc
View File

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

17
tensorflow/c/tf_tensor.cc
View File

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

22
tensorflow/c/tf_tensor.h
View File

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

5
tensorflow/c/tf_tensor_internal.h
View File

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