Skip to content

Commit 6ca2ce1

Browse files
tensorflower-gardenertensorflower-gardener
committed
Allow high-rank Dilate
PiperOrigin-RevId: 935099162
1 parent 839c3b6 commit 6ca2ce1

2 files changed

Lines changed: 48 additions & 67 deletions

File tree

‎tensorflow/lite/kernels/dilate.cc‎

Lines changed: 19 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515
#include <algorithm>
16-
#include <array>
1716
#include <cstdint>
1817
#include <cstring>
1918
#include <vector>
@@ -60,10 +59,6 @@ namespace builtin {
6059
namespace dilate {
6160
namespace {
6261

63-
constexpr size_t kMaxDilateDims = 6;
64-
65-
using Array = std::array<int32_t, kMaxDilateDims>;
66-
6762
// Recursive implementation of the dilation.
6863
//
6964
// This is implemented as a strided copy of the input elements interleaved with
@@ -107,15 +102,14 @@ class DilationRunner {
107102
public:
108103
DilationRunner(const TfLiteIntArray& shape, const int32_t* const dilations,
109104
const char* padding_value, const int element_size)
110-
: size_(shape.size), element_size_(element_size) {
111-
static_assert(sizeof(shape.data[0]) == sizeof(Array::value_type),
112-
"Don't use memcpy here if you change the Array type.");
113-
std::memcpy(shape_.data(), shape.data, size_ * sizeof(shape.data[0]));
114-
static_assert(sizeof(dilations[0]) == sizeof(Array::value_type),
115-
"Don't use memcpy here if you change the Array type.");
116-
std::memcpy(dilations_.data(), dilations, size_ * sizeof(dilations[0]));
117-
105+
: shape_(shape.data, shape.data + shape.size),
106+
dilations_(dilations, dilations + shape.size),
107+
size_(shape.size),
108+
element_size_(element_size) {
118109
MergeTrailingDilations();
110+
input_strides_.resize(size_);
111+
output_strides_.resize(size_);
112+
output_element_sizes_.resize(size_);
119113
ComputeInputStrides();
120114
ComputeOutputStridesAndElementSizes();
121115
FillPaddingValueBuffer(padding_value, element_size);
@@ -124,11 +118,13 @@ class DilationRunner {
124118
int size() const { return size_; }
125119
int element_size() const { return element_size_; }
126120
const char* padding_values() const { return padding_value_buffer_.data(); }
127-
const Array& shape() const { return shape_; }
128-
const Array& dilations() const { return dilations_; }
129-
const Array& input_strides() const { return input_strides_; }
130-
const Array& output_strides() const { return output_strides_; }
131-
const Array& output_element_sizes() const { return output_element_sizes_; }
121+
const std::vector<int32_t>& shape() const { return shape_; }
122+
const std::vector<int32_t>& dilations() const { return dilations_; }
123+
const std::vector<int32_t>& input_strides() const { return input_strides_; }
124+
const std::vector<int32_t>& output_strides() const { return output_strides_; }
125+
const std::vector<int32_t>& output_element_sizes() const {
126+
return output_element_sizes_;
127+
}
132128

133129
void Run(const char* const input, char* const output) const {
134130
DilateImpl(input, output, padding_values(), size(), shape().data(),
@@ -204,11 +200,11 @@ class DilationRunner {
204200
}
205201
}
206202

207-
Array shape_;
208-
Array dilations_;
209-
Array output_strides_;
210-
Array output_element_sizes_;
211-
Array input_strides_;
203+
std::vector<int32_t> shape_;
204+
std::vector<int32_t> dilations_;
205+
std::vector<int32_t> output_strides_;
206+
std::vector<int32_t> output_element_sizes_;
207+
std::vector<int32_t> input_strides_;
212208
// Holds copies of the padding value to memcpy to the output tensor.
213209
std::vector<char> padding_value_buffer_;
214210
int size_;
@@ -263,7 +259,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
263259
const DilationContext ctx(context, node);
264260
TF_LITE_ENSURE(context, ctx.input_tensor->dims != nullptr);
265261
TF_LITE_ENSURE(context, ctx.input_tensor->dims->size > 0);
266-
TF_LITE_ENSURE(context, ctx.input_tensor->dims->size <= kMaxDilateDims);
267262
TF_LITE_ENSURE_EQ(context, ctx.input_tensor->type, ctx.output_tensor->type);
268263
TF_LITE_ENSURE_EQ(context, ctx.input_tensor->type,
269264
ctx.padding_value_tensor->type);

‎tensorflow/lite/kernels/dilate_test.cc‎

Lines changed: 29 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -40,70 +40,42 @@ std::vector<T> DilateReference(const std::vector<T>& input,
4040
const std::vector<int32_t>& shape,
4141
const std::vector<int32_t>& dilations,
4242
const T padding_value) {
43-
constexpr int kMaxDilateDims = 6;
44-
4543
// Compute the output shape.
46-
std::vector<int> output_shape(kMaxDilateDims, 0);
44+
std::vector<int> output_shape(shape.size(), 0);
4745
for (size_t i = 0; i < shape.size(); ++i) {
4846
output_shape[i] = (shape[i] - 1) * dilations[i] + 1;
4947
}
5048

5149
// Compute the input strides.
52-
std::vector<int> strides(kMaxDilateDims, 0);
50+
std::vector<int> strides(shape.size(), 0);
5351
strides[shape.size() - 1] = 1;
5452
for (size_t i = shape.size() - 1; i > 0; --i) {
5553
strides[i - 1] = shape[i] * strides[i];
5654
}
5755

5856
// Compute the output strides.
59-
std::vector<int> output_strides(kMaxDilateDims, 0);
57+
std::vector<int> output_strides(shape.size(), 0);
6058
output_strides[shape.size() - 1] = 1;
6159
for (size_t i = shape.size() - 1; i > 0; --i) {
6260
output_strides[i - 1] = output_shape[i] * output_strides[i];
6361
}
6462

65-
// Copy the dilations to a buffer that can hold the maximum ranks.
66-
std::vector<int> safe_dilations(kMaxDilateDims, 0);
67-
absl::c_copy(dilations, safe_dilations.begin());
68-
69-
// Copy the input shape to a buffer that can hold the maximum ranks.
70-
std::vector<int> safe_input_shape(kMaxDilateDims, 0);
71-
absl::c_copy(shape, safe_input_shape.begin());
72-
7363
// Create a buffer that can hold the output data filled with 0.
7464
std::vector<T> output(
75-
std::accumulate(output_shape.begin(), output_shape.begin() + shape.size(),
76-
1, std::multiplies<>()),
65+
std::accumulate(output_shape.begin(), output_shape.end(), 1,
66+
std::multiplies<>()),
7767
padding_value);
7868

79-
int a = 0;
80-
do {
81-
int b = 0;
82-
do {
83-
int c = 0;
84-
do {
85-
int d = 0;
86-
do {
87-
int e = 0;
88-
do {
89-
int f = 0;
90-
do {
91-
const int i_idx = a * strides[0] + b * strides[1] +
92-
c * strides[2] + d * strides[3] +
93-
e * strides[4] + f * strides[5];
94-
const int o_idx = a * safe_dilations[0] * output_strides[0] +
95-
b * safe_dilations[1] * output_strides[1] +
96-
c * safe_dilations[2] * output_strides[2] +
97-
d * safe_dilations[3] * output_strides[3] +
98-
e * safe_dilations[4] * output_strides[4] +
99-
f * safe_dilations[5] * output_strides[5];
100-
output[o_idx] = input[i_idx];
101-
} while (++f < safe_input_shape[5]);
102-
} while (++e < safe_input_shape[4]);
103-
} while (++d < safe_input_shape[3]);
104-
} while (++c < safe_input_shape[2]);
105-
} while (++b < safe_input_shape[1]);
106-
} while (++a < safe_input_shape[0]);
69+
for (int input_index = 0; input_index < input.size(); ++input_index) {
70+
int remaining_index = input_index;
71+
int output_index = 0;
72+
for (int dim = 0; dim < shape.size(); ++dim) {
73+
const int coordinate = remaining_index / strides[dim];
74+
remaining_index %= strides[dim];
75+
output_index += coordinate * dilations[dim] * output_strides[dim];
76+
}
77+
output[output_index] = input[input_index];
78+
}
10779

10880
return output;
10981
}
@@ -370,5 +342,19 @@ TYPED_TEST(DilateTest, CheckAgainstReferenceImplementation) {
370342
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected));
371343
}
372344

345+
TYPED_TEST(DilateTest, RankSeven) {
346+
auto& model = this->model_;
347+
model.SetInput(/*shape=*/{2, 1, 2, 1, 2, 1, 2});
348+
model.SetDilations(/*dilations=*/{2, 1, 2, 1, 1, 1, 2});
349+
model.SetPaddingValue(-1);
350+
351+
const auto expected =
352+
DilateReference(model.GetInput(), model.GetInputShape(),
353+
model.GetDilations(), model.GetPaddingValue());
354+
EXPECT_EQ(model.BuildAndInvoke(), kTfLiteOk);
355+
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({3, 1, 3, 1, 2, 1, 3}));
356+
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected));
357+
}
358+
373359
} // namespace
374360
} // namespace tflite

0 commit comments

Comments
 (0)