1/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15#include <stddef.h>
16#include <stdint.h>
17
18#include "tensorflow/lite/c/common.h"
19#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
20#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
21#include "tensorflow/lite/kernels/internal/tensor.h"
22#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
23#include "tensorflow/lite/kernels/kernel_util.h"
24
25// TODO(b/117523611): We should factor out a binary_op and put binary ops there.
26namespace tflite {
27namespace ops {
28namespace builtin {
29namespace floor_mod {
30namespace {
31
32// Input/output tensor index.
33constexpr int kInputTensor1 = 0;
34constexpr int kInputTensor2 = 1;
35constexpr int kOutputTensor = 0;
36
37// Op data for floor_mod op.
38struct OpData {
39 bool requires_broadcast;
40};
41
42// TODO(b/117912880): Support quantization.
43
44void* Init(TfLiteContext* context, const char* buffer, size_t length) {
45 auto* data = new OpData;
46 data->requires_broadcast = false;
47 return data;
48}
49
50void Free(TfLiteContext* context, void* buffer) {
51 delete reinterpret_cast<OpData*>(buffer);
52}
53
54TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
55 TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
56 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
57
58 // Reinterprete the opaque data provided by user.
59 OpData* data = reinterpret_cast<OpData*>(node->user_data);
60
61 const TfLiteTensor* input1;
62 TF_LITE_ENSURE_OK(context,
63 GetInputSafe(context, node, kInputTensor1, &input1));
64 const TfLiteTensor* input2;
65 TF_LITE_ENSURE_OK(context,
66 GetInputSafe(context, node, kInputTensor2, &input2));
67 TfLiteTensor* output;
68 TF_LITE_ENSURE_OK(context,
69 GetOutputSafe(context, node, kOutputTensor, &output));
70
71 TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
72
73 const TfLiteType type = input1->type;
74 if (type != kTfLiteInt32 && type != kTfLiteFloat32 && type != kTfLiteInt64) {
75 TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by floor_mod.",
76 TfLiteTypeGetName(type));
77 return kTfLiteError;
78 }
79 output->type = type;
80
81 data->requires_broadcast = !HaveSameShapes(input1, input2);
82
83 TfLiteIntArray* output_size = nullptr;
84 if (data->requires_broadcast) {
85 TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
86 context, input1, input2, &output_size));
87 } else {
88 output_size = TfLiteIntArrayCopy(input1->dims);
89 }
90
91 return context->ResizeTensor(context, output, output_size);
92}
93
94template <typename T>
95TfLiteStatus EvalImpl(TfLiteContext* context, bool requires_broadcast,
96 const TfLiteTensor* input1, const TfLiteTensor* input2,
97 TfLiteTensor* output) {
98 const T* denominator_data = GetTensorData<T>(input2);
99
100 if (input2->type == kTfLiteInt32 || input2->type == kTfLiteInt64) {
101 // Validate the denominator only for integer.
102 const int num_elements = NumElements(input2);
103 for (int i = 0; i < num_elements; ++i) {
104 if (denominator_data[i] == 0) {
105 TF_LITE_KERNEL_LOG(context, "Division by 0");
106 return kTfLiteError;
107 }
108 }
109 }
110 if (requires_broadcast) {
111 reference_ops::BroadcastBinaryFunction4DSlow<T, T, T>(
112 GetTensorShape(input1), GetTensorData<T>(input1),
113 GetTensorShape(input2), denominator_data, GetTensorShape(output),
114 GetTensorData<T>(output), reference_ops::FloorMod<T>);
115 } else {
116 reference_ops::BinaryFunction<T, T, T>(
117 GetTensorShape(input1), GetTensorData<T>(input1),
118 GetTensorShape(input2), GetTensorData<T>(input2),
119 GetTensorShape(output), GetTensorData<T>(output),
120 reference_ops::FloorMod<T>);
121 }
122
123 return kTfLiteOk;
124}
125
126TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
127 OpData* data = reinterpret_cast<OpData*>(node->user_data);
128
129 const TfLiteTensor* input1;
130 TF_LITE_ENSURE_OK(context,
131 GetInputSafe(context, node, kInputTensor1, &input1));
132 const TfLiteTensor* input2;
133 TF_LITE_ENSURE_OK(context,
134 GetInputSafe(context, node, kInputTensor2, &input2));
135 TfLiteTensor* output;
136 TF_LITE_ENSURE_OK(context,
137 GetOutputSafe(context, node, kOutputTensor, &output));
138
139 switch (input1->type) {
140 case kTfLiteInt32: {
141 return EvalImpl<int32_t>(context, data->requires_broadcast, input1,
142 input2, output);
143 }
144 case kTfLiteInt64: {
145 return EvalImpl<int64_t>(context, data->requires_broadcast, input1,
146 input2, output);
147 }
148 case kTfLiteFloat32: {
149 return EvalImpl<float>(context, data->requires_broadcast, input1, input2,
150 output);
151 }
152 default: {
153 TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by floor_mod.",
154 TfLiteTypeGetName(input1->type));
155 return kTfLiteError;
156 }
157 }
158}
159
160} // namespace
161} // namespace floor_mod
162
163TfLiteRegistration* Register_FLOOR_MOD() {
164 // Init, Free, Prepare, Eval are satisfying the Interface required by
165 // TfLiteRegistration.
166 static TfLiteRegistration r = {floor_mod::Init, floor_mod::Free,
167 floor_mod::Prepare, floor_mod::Eval};
168 return &r;
169}
170
171} // namespace builtin
172} // namespace ops
173} // namespace tflite
174