Update kernel evals to use new kernel signatures.
PiperOrigin-RevId: 214763814
This commit is contained in:
parent
19d8963bc0
commit
abf2635620
tensorflow/contrib/lite/kernels
@ -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:
|
||||
|
@ -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);
|
||||
|
@ -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_;
|
||||
};
|
||||
|
@ -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.
|
||||
|
@ -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) { \
|
||||
|
@ -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: {
|
||||
|
@ -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:
|
||||
|
@ -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:
|
||||
|
@ -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) {
|
||||
|
Loading…
Reference in New Issue
Block a user