slice_utils.h 5.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2021 PaddlePaddle 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. */

#pragma once
H
hong 已提交
16
#include <paddle/phi/core/ddim.h>
17

18 19 20
#include <string>
#include <vector>

H
hong 已提交
21 22 23
namespace phi {

namespace funcs {
24 25

template <typename T = int64_t>
H
hong 已提交
26
inline void CheckAndUpdateSliceAttrs(const DDim in_dims,
27 28 29 30 31 32 33
                                     const std::vector<T>& axes,
                                     std::vector<T>* starts,
                                     std::vector<T>* ends,
                                     std::vector<int64_t>* steps = nullptr,
                                     std::vector<T>* infer_flags = nullptr) {
  for (size_t i = 0; i < axes.size(); ++i) {
    T axis = axes[i];
34
    PADDLE_ENFORCE_LT(
H
hong 已提交
35 36 37
        axis,
        in_dims.size(),
        phi::errors::InvalidArgument(
38 39
            "The axis value should be less than the rank of input, "
            "but received axes[%d] = %d, rank of input is %d.",
H
hong 已提交
40 41 42
            i,
            axis,
            in_dims.size()));
43 44 45 46 47

    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
      continue;
    }

48 49 50 51 52
    T dim_value = in_dims[axis];

    if (dim_value > 0) {
      T step = steps == nullptr ? 1 : (*steps)[i];
      PADDLE_ENFORCE_NE(
H
hong 已提交
53 54 55 56
          step,
          0,
          phi::errors::InvalidArgument(
              "Step should not be 0, but received step = %d.", step));
57

58 59 60 61 62 63 64
      T start = (*starts)[i] < 0 ? ((*starts)[i] + dim_value) : (*starts)[i];
      start = std::max(start, static_cast<T>(0));

      T end =
          0 < step && (*ends)[i] < 0 ? ((*ends)[i] + dim_value) : (*ends)[i];
      end = std::min(end, dim_value);

65 66 67
      if (step > 0) {
        start = std::min(start, dim_value);
        end = std::max(end, static_cast<T>(0));
68
        PADDLE_ENFORCE_GE(
H
hong 已提交
69 70 71
            end,
            start,
            phi::errors::InvalidArgument(
72 73
                "When step > 0, end should be greater than start, but "
                "received end = %d, start = %d.",
H
hong 已提交
74 75
                end,
                start));
76 77 78 79 80
      } else {
        // NOTE(liym27): When step < 0, start should less and equal to
        // dim_value-1
        // "end is -1" means contain the 0-th element of this axis.
        start = std::min(start, dim_value - 1);
Z
zyfncg 已提交
81 82 83
        if (end < -1) {
          end += dim_value;
        }
84
        end = std::max(end, static_cast<T>(-1));
85
        PADDLE_ENFORCE_GE(
H
hong 已提交
86 87 88
            start,
            end,
            phi::errors::InvalidArgument(
89 90
                "When step < 0, start should be greater than end, but "
                "received start = %d, end = %d.",
H
hong 已提交
91 92
                start,
                end));
93 94 95 96
      }

      (*starts)[i] = start;
      (*ends)[i] = end;
97 98 99
    } else if (dim_value == 0) {
      (*starts)[i] = 0;
      (*ends)[i] = 0;
100 101 102 103 104
    }
  }
}

template <typename T = int64_t>
H
hong 已提交
105 106 107 108 109 110 111
inline phi::DDim GetSliceDims(const phi::DDim in_dims,
                              const std::vector<T>& axes,
                              const std::vector<T>& starts,
                              const std::vector<T>& ends,
                              std::vector<T>* steps = nullptr,
                              std::vector<T>* infer_flags = nullptr) {
  phi::DDim slice_dims(in_dims);
112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133

  for (size_t i = 0; i < axes.size(); ++i) {
    T axis = axes[i];
    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
      slice_dims[axis] = -1;
      continue;
    }

    T start = starts[i];
    T end = ends[i];
    T step = steps == nullptr ? 1 : (*steps)[i];

    if (step > 0) {
      slice_dims[axis] = (end - start + step - 1) / step;
    } else {
      slice_dims[axis] = (end - start + step + 1) / step;
    }
  }
  return slice_dims;
}

template <typename T = int64_t>
H
hong 已提交
134 135 136 137
inline DDim GetDecreasedDims(const DDim slice_dims,
                             const std::vector<T>& decrease_axes,
                             std::vector<T>* infer_flags = nullptr) {
  DDim decreased_dims(slice_dims);
138
  std::vector<uint8_t> decrease_flag(slice_dims.size(), 0);
139 140 141
  if (decrease_axes.size() > 0) {
    for (size_t i = 0; i < decrease_axes.size(); ++i) {
      T axis = decrease_axes[i];
142
      decrease_flag[axis] = 1;
143
      if (infer_flags && (*infer_flags)[i] != -1) {
H
hong 已提交
144 145 146
        PADDLE_ENFORCE_EQ(decreased_dims[axis],
                          1,
                          phi::errors::InvalidArgument(
147 148
                              "Decrease dim should be 1, but now received %d",
                              decreased_dims[axis]));
149 150 151 152 153
      }
    }

    std::vector<T> new_shape;
    for (int i = 0; i < decreased_dims.size(); ++i) {
154
      if (decrease_flag[i] == 0) {
155 156 157 158 159 160 161 162 163 164
        new_shape.push_back(decreased_dims[i]);
      }
    }

    // NOTE(liym27): Paddle does not support that the rank of Tensor is 0, and
    // uses [1] instead.
    if (new_shape.size() == 0) {
      new_shape.push_back(1);
    }

165
    decreased_dims = phi::make_ddim(new_shape);
166 167 168 169
  }
  return decreased_dims;
}

H
hong 已提交
170 171
}  // namespace funcs
}  // namespace phi