Update kernel evals to use new kernel signatures.

PiperOrigin-RevId: 214763814
This commit is contained in:
A. Unique TensorFlower 2018-09-27 06:12:59 -07:00 committed by TensorFlower Gardener
parent 19d8963bc0
commit abf2635620
9 changed files with 92 additions and 111 deletions

View File

@ -100,20 +100,31 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
// allocate and populate these during Prepare().
// TODO(ycling): Activation function parameter is ignored. For now we dont have
// a model with a Concatenation with fused activation function.
#define TF_LITE_CONCATENATION(type, scalar) \
VectorOfTensors<scalar> all_inputs(*context, *node->inputs); \
type::Concatenation<FusedActivationFunctionType::kNone, scalar>( \
RemapDim(NumDimensions(output), axis), all_inputs.data(), \
all_inputs.dims(), node->inputs->size, GetTensorData<scalar>(output), \
GetTensorDims(output))
#define TF_LITE_CONCATENATION(type, scalar) \
{ \
VectorOfTensors<scalar> all_inputs(*context, *node->inputs); \
tflite::ConcatenationParams op_params; \
op_params.axis = axis; \
op_params.inputs_count = node->inputs->size; \
type::Concatenation(op_params, all_inputs.shapes(), all_inputs.data(), \
GetTensorShape(output), \
GetTensorData<scalar>(output)); \
}
#define TF_LITE_CONCATENATION_QUANTIZED(type) \
VectorOfQuantizedTensors all_inputs(*context, *node->inputs); \
type::Concatenation( \
RemapDim(NumDimensions(output), axis), all_inputs.data(), \
all_inputs.dims(), all_inputs.zero_point(), all_inputs.scale(), \
node->inputs->size, GetTensorData<uint8>(output), GetTensorDims(output), \
output->params.zero_point, output->params.scale)
#define TF_LITE_CONCATENATION_QUANTIZED(type) \
{ \
VectorOfQuantizedTensors all_inputs(*context, *node->inputs); \
tflite::ConcatenationParams op_params; \
op_params.axis = axis; \
op_params.input_zeropoint = all_inputs.zero_point(); \
op_params.input_scale = all_inputs.scale(); \
op_params.inputs_count = node->inputs->size; \
op_params.output_zeropoint = output->params.zero_point; \
op_params.output_scale = output->params.scale; \
type::ConcatenationWithScaling(op_params, all_inputs.shapes(), \
all_inputs.data(), GetTensorShape(output), \
GetTensorData<uint8>(output)); \
}
switch (output->type) { // Already know in/outtypes are same.
case kTfLiteFloat32:

View File

@ -84,11 +84,15 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* positions = GetInput(context, node, kInputPositions);
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
const int input_rank = NumDimensions(input);
#define TF_LITE_GATHER(data_type, index_type) \
optimized_ops::Gather( \
GetTensorData<data_type>(input), GetTensorDims(input), input_rank, \
GetTensorData<index_type>(positions), GetTensorDims(positions), \
GetTensorData<data_type>(output), GetTensorDims(output));
#define TF_LITE_GATHER(data_type, index_type) \
{ \
tflite::GatherParams op_params; \
op_params.input_rank = input_rank; \
optimized_ops::Gather( \
op_params, GetTensorShape(input), GetTensorData<data_type>(input), \
GetTensorShape(positions), GetTensorData<index_type>(positions), \
GetTensorShape(output), GetTensorData<data_type>(output)); \
}
switch (input->type) {
case kTfLiteFloat32:
TF_LITE_GATHER(float, int32_t);

View File

@ -56,23 +56,19 @@ class VectorOfTensors {
int num_tensors = tensor_list.size;
all_data_.reserve(num_tensors);
all_dims_.reserve(num_tensors);
all_dims_ptr_.reserve(num_tensors);
all_shape_.reserve(num_tensors);
all_shape_ptr_.reserve(num_tensors);
for (int i = 0; i < num_tensors; ++i) {
TfLiteTensor* t = &context.tensors[tensor_list.data[i]];
all_data_.push_back(GetTensorData<T>(t));
all_dims_.push_back(GetTensorDims(t));
all_shape_.push_back(GetTensorShape(t));
}
// Taking the pointer from inside a std::vector is only OK if the vector is
// never modified, so we populate all_dims in the previous loop and then we
// never modified, so we populate all_shape in the previous loop and then we
// are free to grab iterators here.
for (int i = 0; i < num_tensors; ++i) {
all_dims_ptr_.push_back(&all_dims_[i]);
all_shape_ptr_.push_back(&all_shape_[i]);
}
}
@ -82,12 +78,6 @@ class VectorOfTensors {
// f[0][1] is the second element of the first tensor.
T* const* data() const { return all_data_.data(); }
// Return a pointer the dim pointers of all tensors in the list. For
// example:
// const Dims<4>* const* d = v.dims();
// dims[1] are the dimensions of the second tensor in the list.
const Dims<4>* const* dims() const { return all_dims_ptr_.data(); }
// Return a pointer the shape pointers of all tensors in the list. For
// example:
// const RuntimeShape* const* d = v.dims();
@ -96,8 +86,6 @@ class VectorOfTensors {
private:
std::vector<T*> all_data_;
std::vector<Dims<4>> all_dims_;
std::vector<Dims<4>*> all_dims_ptr_;
std::vector<RuntimeShape> all_shape_;
std::vector<RuntimeShape*> all_shape_ptr_;
};

View File

@ -86,10 +86,6 @@ inline const bool* GetTensorData(const TfLiteTensor* tensor) {
return tensor != nullptr ? tensor->data.b : nullptr;
}
inline int RemapDim(int max_dimensions, int d) {
return max_dimensions - d - 1;
}
// TODO(ahentz): the implementations in kernels/internal/ take a Dims<4> object
// even if the original tensors were not 4D. We should consider rewriting them
// to take a more generic 'shape' object.

View File

@ -70,12 +70,12 @@ TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) {
bool is_rank_one = !HaveSameShapes(input_condition, input_x);
#define TF_LITE_SELECT(type, op) \
reference_ops::op(GetTensorData<bool>(input_condition), \
GetTensorDims(input_condition), \
GetTensorData<type>(input_x), GetTensorDims(input_x), \
GetTensorData<type>(input_y), GetTensorDims(input_y), \
GetTensorData<type>(output), GetTensorDims(output));
#define TF_LITE_SELECT(type, op) \
reference_ops::op(GetTensorShape(input_condition), \
GetTensorData<bool>(input_condition), \
GetTensorShape(input_x), GetTensorData<type>(input_x), \
GetTensorShape(input_y), GetTensorData<type>(input_y), \
GetTensorShape(output), GetTensorData<type>(output));
#define TF_LITE_SWITCH(type, op) \
switch (type) { \

View File

@ -109,25 +109,24 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
if (axis_value < 0) {
axis_value += NumDimensions(op_context.input);
}
axis_value = RemapDim(NumDimensions(op_context.input), axis_value);
// TODO(ahentz): Our usage of VectorOfTensors could be optimized by
// calculating it in Prepare, unless we defer shape calculation.
// TODO(ahentz): We can improve the optimized_ops version to handle other
// cases too.
#define TF_LITE_SPLIT(scalar) \
VectorOfTensors<scalar> all_outputs(*context, *node->outputs); \
if (axis_value == NumDimensions(op_context.input)) { \
optimized_ops::TensorFlowSplit<FusedActivationFunctionType::kNone, \
scalar>( \
GetTensorData<scalar>(op_context.input), \
GetTensorDims(op_context.input), NumOutputs(node), all_outputs.data(), \
all_outputs.dims()); \
} else { \
reference_ops::TensorFlowSplit<scalar>( \
GetTensorData<scalar>(op_context.input), \
GetTensorDims(op_context.input), axis_value, NumOutputs(node), \
all_outputs.data(), all_outputs.dims()); \
#define TF_LITE_SPLIT(scalar) \
VectorOfTensors<scalar> all_outputs(*context, *node->outputs); \
tflite::SplitParams op_params; \
op_params.num_split = NumOutputs(node); \
op_params.axis = axis_value; \
if (axis_value == 0) { \
optimized_ops::Split(op_params, GetTensorShape(op_context.input), \
GetTensorData<scalar>(op_context.input), \
all_outputs.shapes(), all_outputs.data()); \
} else { \
reference_ops::Split(op_params, GetTensorShape(op_context.input), \
GetTensorData<scalar>(op_context.input), \
all_outputs.shapes(), all_outputs.data()); \
}
switch (op_context.input->type) {
case kTfLiteFloat32: {

View File

@ -57,17 +57,6 @@ struct StridedSliceContext {
int dims;
};
// Reverse order of bits in the mask to match the expected order in kernel
inline int ReverseMaskBits(int mask, int num_dimensions) {
int out = 0;
for (int dim = 0; dim < num_dimensions; dim++) {
out <<= 1;
out += (mask & 1);
mask >>= 1;
}
return out;
}
// This Op only supports 1-4D cases and since we use the reference 4D
// implementation, the 1-3D tensors are mapped to 4D.
const int kMaxDim = 4;
@ -198,30 +187,31 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
std::vector<int32_t> stops;
std::vector<int32_t> strides;
for (int idx = op_context.dims - 1; idx >= 0; --idx) {
starts.emplace_back(GetTensorData<int32_t>(op_context.begin)[idx]);
stops.emplace_back(GetTensorData<int32_t>(op_context.end)[idx]);
strides.emplace_back(GetTensorData<int32_t>(op_context.strides)[idx]);
}
for (int i = op_context.dims; i < kMaxDim; i++) {
starts.emplace_back(0);
stops.emplace_back(1);
strides.emplace_back(1);
}
int begin_mask =
ReverseMaskBits(op_context.params->begin_mask, op_context.dims);
int end_mask = ReverseMaskBits(op_context.params->end_mask, op_context.dims);
int shrink_axis_mask =
ReverseMaskBits(op_context.params->shrink_axis_mask, op_context.dims);
for (int idx = 0; idx < op_context.dims; ++idx) {
starts.emplace_back(GetTensorData<int32_t>(op_context.begin)[idx]);
stops.emplace_back(GetTensorData<int32_t>(op_context.end)[idx]);
strides.emplace_back(GetTensorData<int32_t>(op_context.strides)[idx]);
}
#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \
kernel_type::StridedSlice( \
GetTensorData<data_type>(op_context.input), \
GetTensorDims(op_context.input), begin_mask, end_mask, shrink_axis_mask, \
starts, stops, strides, GetTensorData<data_type>(op_context.output), \
GetTensorDims(op_context.output))
int begin_mask = op_context.params->begin_mask << (4 - op_context.dims);
int end_mask = op_context.params->end_mask << (4 - op_context.dims);
int shrink_axis_mask = op_context.params->shrink_axis_mask
<< (4 - op_context.dims);
TF_LITE_ENSURE_EQ(context, starts.size(), 4);
auto op_params = ::tflite::strided_slice::BuildStridedSliceParams(
begin_mask, end_mask, shrink_axis_mask, starts, stops, strides);
#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \
kernel_type::StridedSlice(op_params, GetTensorShape(op_context.input), \
GetTensorData<data_type>(op_context.input), \
GetTensorShape(op_context.output), \
GetTensorData<data_type>(op_context.output))
switch (op_context.input->type) {
case kTfLiteFloat32:

View File

@ -92,26 +92,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
}
// Reverse the permuted axes and convert to 4D due to the way Dims are
// constructed in GetTensorDims.
const int* perm_data = GetTensorData<int32_t>(op_context.perm);
const int size = op_context.perm->dims->data[0];
const int kOutputDimensionNum = 4;
int reversed_perm[kOutputDimensionNum];
for (int output_k = 0, input_k = size - 1; output_k < size;
++output_k, --input_k) {
reversed_perm[output_k] = size - perm_data[input_k] - 1;
}
for (int k = size; k < kOutputDimensionNum; ++k) {
reversed_perm[k] = k;
TransposeParams params;
params.perm_count = size;
for (int i = 0; i < size; ++i) {
params.perm[i] = perm_data[i];
}
#define TF_LITE_TRANSPOSE(type, scalar) \
type::Transpose(GetTensorData<scalar>(op_context.input), \
GetTensorDims(op_context.input), \
GetTensorData<scalar>(op_context.output), \
GetTensorDims(op_context.output), reversed_perm)
type::Transpose(params, GetTensorShape(op_context.input), \
GetTensorData<scalar>(op_context.input), \
GetTensorShape(op_context.output), \
GetTensorData<scalar>(op_context.output))
switch (op_context.input->type) {
case kTfLiteFloat32:

View File

@ -51,21 +51,21 @@ void RunTestPermutation(const std::vector<int>& shape,
reversed_perms[k] = k;
}
// Make input and output dims (i.e. reversed shape and dest_shape).
Dims<4> input_dims = GetTensorDims(shape);
Dims<4> output_dims;
for (int i = 0; i < 4; i++) {
output_dims.sizes[i] = input_dims.sizes[reversed_perms[i]];
}
output_dims.strides[0] = 1;
for (int k = 1; k < 4; k++) {
output_dims.strides[k] =
output_dims.strides[k - 1] * output_dims.sizes[k - 1];
// Make input and output shapes.
const RuntimeShape input_shape = GetTensorShape(shape);
RuntimeShape output_shape(perms.size());
for (int i = 0; i < perms.size(); i++) {
output_shape.SetDim(i, input_shape.Dims(perms[i]));
}
reference_ops::Transpose<float>(input.data(), input_dims,
input_transposed->data(), output_dims,
reversed_perms);
TransposeParams params;
params.perm_count = perms.size();
for (int i = 0; i < perms.size(); ++i) {
params.perm[i] = perms[i];
}
reference_ops::Transpose<float>(params, input_shape, input.data(),
output_shape, input_transposed->data());
}
TEST(TransposeTest, TestRefOps1D) {