chromium/third_party/tflite/src/tensorflow/lite/kernels/dilate.cc

/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include <vector>

#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/util.h"

// This file implements a dilation operation on a tensor.
//
// The dilation operation scatters the elements of its input into a new tensor
// according to a dilation factor for each dimension. The new tensor elements
// are initialized to 0.
//
// This operation can also be seen as adding interior padding to the tensor. In
// that case, `interior padding size = dilation factor - 1`.
//
// For instance:
//
//                        1 2 3
// A is a 3x3 tensor. A = 4 5 6
//                        7 8 9
//
// We apply a dilation of 2x3.
//
//                         1 0 0 2 0 0 3
//                         0 0 0 0 0 0 0
// B = dilate(A, [2, 3]) = 4 0 0 5 0 0 6
//                         0 0 0 0 0 0 0
//                         7 0 0 8 0 0 9
//
// More rigorously:
// - Let [s0, ..., sN] be the shape of A.
// - Let [d0, ..., dN] be the dilation factors.
//
// - The shape of B is [(s0 - 1) * d0 + 1, ..., (sN - 1) * dN + 1].
// - B(i0, ..., iN) = ┌ A(i0 / d0, ..., iN / dN)   if iX % dX == 0 for all X
//                    └ 0 otherwise.

namespace tflite {
namespace ops {
namespace builtin {
namespace dilate {
namespace {

constexpr size_t kMaxDilateDims = …;

Array;

// Recursive implementation of the dilation.
//
// This is implemented as a strided copy of the input elements interleaved with
// calls to memset to zero out the padding elements.
void DilateImpl(const char* input, char* output,
                const char* const padding_values, const int32_t size,
                const int32_t* const shape, const int32_t* const input_strides,
                const int32_t* const output_strides,
                const int32_t* const output_element_sizes, size_t depth = 0) { … }

// Prepares the data needed by the dilation actual implementation.
//
// This class also has an optimization pass to reduce the number of calls to
// memcpy in the implementation.
class DilationRunner { … };

// Holds the tensors and operation context for convenience.
struct DilationContext { … };

// Computes the new length of a dimension given its dilation factor.
int DilateDim(int dim, int dilate_factor) { … }

// Computes the output tensor shape and resizes it.
TfLiteStatus SetupOutputTensor(const DilationContext& ctx) { … }

// Prepares the dilate operation.
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }

// Runs the dilate operation.
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { … }

}  // namespace
}  // namespace dilate

TfLiteRegistration* Register_DILATE() { … }

}  // namespace builtin
}  // namespace ops
}  // namespace tflite