stablehlo_reduce_window.cc | Explore in Territory

/* 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 <cassert>
#include <cstdint>
#include <cstring>
#include <functional>
#include <limits>
#include <memory>
#include <type_traits>
#include <vector>

#include "tensorflow/lite/array.h"
#include "tensorflow/lite/builtin_ops.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/core/c/builtin_op_data.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/core/subgraph.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/util.h"

namespace tflite {
namespace ops {
namespace builtin {

namespace {
constexpr int32_t kMaxReduceWindowRank = …;

// Reccursive implementation of a strided copy of a tensor.
void StridedCopy(const int rank, const char* input, const int64_t* input_shape,
                 const int64_t* input_strides, char* output,
                 const int64_t* output_strides, const int64_t element_size,
                 const int depth) { … }

}  // namespace

namespace dilate {
namespace {

const int64_t kTFLiteDefaultBaseDilation[kMaxReduceWindowRank] = …;

// Computes and holds the parameters that can be precomputed for the dilation
// operation.
struct DilateData { … };

// Dilates the input tensor following the parameters held in the given context.
//
// 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.
void Dilate(const DilateData& ctx, const char* input, const char* init_value,
            char* output) { … }

}  // namespace
}  // namespace dilate

namespace pad {
namespace {

const int64_t kTFLiteDefaultPadding[kMaxReduceWindowRank] = …;

// Computes and holds the parameters that can be precomputed for the padding
// operation. Note that StableHLO padding treats negative values as cropping.
struct PadCropData { … };

// Pads and crops the input tensor following the parameters held in the given
// context.
//
// The StableHLO padding algorithm uses negative values to denote cropping.
void PadCrop(const PadCropData& ctx, const char* input, const char* init_value,
             char* output) { … }

}  // namespace
}  // namespace pad

namespace reduce_window {
namespace {

// Reduces the elements of a tensor viewed through a strided window.
//
// This applies a reduction to a tensor by skipping over elements that are not
// in the window defined by the given shape and strides. The window is reduced
// to one element.
//
// The shape is the shape of the window. The strides are based on the actual
// tensor and the distance between window elements, counted in elements.
// Sparse windows are possible.
//
// For instance: the following window has a [2, 2] shape and [8, 3] strides.
//
// ┌──┐     ┌──┐
// │ 1│ 2  3│ 4│
// └──┘     └──┘
//   5  6  7  8    is reduced to 1 + 4 + 9 + 12 = 26
// ┌──┐     ┌──┐
// │ 9│10 11│12│
// └──┘     └──┘
//  13 14 15 16
//
// This is a recursive implementation of the strided reduction.
template <class Op, class Type>
void StridedReduce(const Type* input, const int64_t* const shape,
                   const int64_t* const strides, Type& accu, const int rank,
                   const int depth) { … }

// Recursively computes strided reductions using a sliding window over the
// given tensor.
//
// The window is defined using a shape and a dilation. The shape defines the
// elements that the window will let the reduction *see*. The dilation defines
// the step between window elements.
//
// For instance: the following window has a [2, 2] shape and [2, 3] dilations.
//
//    3
// ┌────┐
// ┌─┐   ┌─┐
// │X│X X│X│┐
// └─┘   └─┘│2
//  X X X X ┘
// ┌─┐   ┌─┐
// │X│X X│X│
// └─┘   └─┘
template <class Op, class Type>
void ReduceWindowImpl(const Type* input, Type* output,
                      const int64_t* const output_shape,
                      const int64_t* const output_strides,
                      const int64_t* const window_offset_strides,
                      const int64_t* const window_shape,
                      const int64_t* const window_reduce_strides,
                      const Type init, const int rank, const int depth) { … }

// Computes and holds the parameters that can be precomputed for the dilation
// operation.
struct ReduceWindowData { … };

template <class Op, class Type>
void ReduceWindow(const ReduceWindowData& ctx, const Type* const input,
                  const Type init, Type* output) { … }

}  // namespace
}  // namespace reduce_window

/// Operator implementation

namespace reduce_window_op {
namespace {

// Holds the data needed throughout the node lifetime.
struct NodeData { … };

// Holds the operation data. This is extended by the StablehloData and the
// TFLiteData classes.
//
// There are two available semantics for this op implementation.
//
// - StablehloData, that models the STABLEHLO_REDUCE_WINDOW op.
// - TFLiteData, that models the DEPRECATED initial REDUCE_WINDOW op.
struct OpData { … };

// Speciliazes OpData for the STABLEHLO_REDUCE_WINDOW operation.
struct StablehloData : public OpData { … };

// Specializes OpData for the REDUCE_WINDOW operation.
struct TFLiteData : public OpData { … };

// Applies the sub-ops that are needed to compute the whole
// [STABLEHLO_]REDUCE_WINDOW op.
//
// The ops that aren't needed are skipped.
template <class Op, class Type>
void PadCropReduceWindow(const OpData& op_ctx) { … }

// Dispatches to the template implementation according to the tensor type.
template <class Op>
TfLiteStatus DispatchReduceWindowType(OpData& ctx) { … }

struct Max { … };

struct Min { … };

// Dispatches to the template instanciation according to the reduction body.
TfLiteStatus DispatchReduceWindowBody(OpData& ctx) { … }

// Initializes the node's user data when the STABLEHLO_REDUCE_WINDOW sematic is
// used.
void* StablehloInit(TfLiteContext* context, const char* options,
                    size_t options_len) { … }

void* TFLiteInit(TfLiteContext* context, const char* options,
                 size_t options_len) { … }

// Frees the node's user data when the STABLEHLO_REDUCE_WINDOW sematic is used.
void Free(TfLiteContext* context, void* node_data) { … }

template <class Semantic>
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }

template <class Semantic>
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { … }

}  // namespace
}  // namespace reduce_window_op

TfLiteRegistration* Register_STABLEHLO_REDUCE_WINDOW() { … }

TfLiteRegistration* Register_REDUCE_WINDOW() { … }

}  // namespace builtin
}  // namespace ops
}  // namespace tflite
chromium/third_party/tflite/src/tensorflow/lite/kernels/stablehlo_reduce_window.cc
