infermeta_utils.h 12.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2022 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

#include <string>
18 19
#include <typeindex>
#include <typeinfo>
20 21
#include <utility>

22 23 24 25 26 27
#include "paddle/phi/common/scalar.h"
#include "paddle/phi/common/scalar_array.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/core/macros.h"
#include "paddle/phi/core/meta_tensor.h"
#include "paddle/phi/core/type_defs.h"
28
#include "paddle/utils/any.h"
29
#include "paddle/utils/flat_hash_map.h"
30 31
#include "paddle/utils/small_vector.h"

32
namespace phi {
33 34 35 36 37 38 39

class InferMetaContext {
 public:
  InferMetaContext() = default;
  explicit InferMetaContext(MetaConfig config) : config_(config) {}

  void SetMetaConfig(MetaConfig config);
40 41
  void EmplaceBackInput(std::shared_ptr<phi::MetaTensor> input);
  void EmplaceBackOutput(std::shared_ptr<phi::MetaTensor> output);
42 43 44
  void EmplaceBackAttr(paddle::any attr);

  void EmplaceBackInputs(
45
      paddle::SmallVector<std::shared_ptr<phi::MetaTensor>> inputs);
46
  void EmplaceBackOutputs(
47
      paddle::SmallVector<std::shared_ptr<phi::MetaTensor>> outputs);
48 49 50 51 52 53

  const std::pair<int, int>& InputRangeAt(size_t idx) const;
  const std::pair<int, int>& OutputRangeAt(size_t idx) const;

  const MetaConfig& GetMetaConfig() const;
  const MetaTensor& InputAt(size_t idx) const;
54 55 56

  paddle::optional<const phi::MetaTensor&> OptionalInputAt(size_t idx) const;

57
  std::vector<MetaTensor> InputsBetween(size_t start, size_t end) const;
58
  MetaTensor* MutableOutputAt(size_t idx);
59
  std::vector<MetaTensor> MutableOutputBetween(size_t start, size_t end);
60 61 62 63 64

  template <typename AttrType>
  AttrType AttrAt(size_t idx) {
    try {
      return paddle::any_cast<AttrType>(attrs_.at(idx));
65
    } catch (paddle::bad_any_cast& e) {
66
      PADDLE_THROW(phi::errors::InvalidArgument(
67 68 69 70
          "Attribute cast error in InferMeta Context, the expected attribute "
          "type is `%s`, but actual attribute type is `%s`.",
          std::type_index(typeid(AttrType)).name(),
          std::type_index(attrs_.at(idx).type()).name()));
71 72 73 74 75 76 77 78 79 80
    }
  }

 private:
  MetaConfig config_;

  // NOTE(chenweihang): Because the MetaTensor is a base class, and MetaTensor
  // objects are all created in each round, so we have to use smart pointer
  // here, maybe we can implemented a new InferMetaContext and a series utils
  // specifically for fluid to avoid using shared_ptr
81 82
  paddle::SmallVector<std::shared_ptr<phi::MetaTensor>> inputs_;
  paddle::SmallVector<std::shared_ptr<phi::MetaTensor>> outputs_;
83 84 85 86 87 88 89
  paddle::SmallVector<paddle::any> attrs_;

  paddle::SmallVector<std::pair<int, int>> input_range_;
  paddle::SmallVector<std::pair<int, int>> output_range_;
};

#define PT_INFER_META(...) \
90
  ::phi::InferMetaFnImpl<decltype(&__VA_ARGS__), &__VA_ARGS__>::Call
91 92 93 94 95 96 97 98 99 100

#define PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(attr_type)           \
  template <typename... Tail>                                                  \
  struct InferMetaFnCallHelper<attr_type, Tail...> {                           \
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs> \
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {          \
      static_assert(out_idx == 0,                                              \
                    "InferMeta's Attributes should appear before Outputs.");   \
      attr_type arg = ctx->AttrAt<attr_type>(attr_idx);                        \
      InferMetaFnCallHelper<                                                   \
101 102
          Tail...>::template Call<in_idx, attr_idx + 1, out_idx>(ctx,          \
                                                                 pargs...,     \
103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
                                                                 arg);         \
    }                                                                          \
  }

template <typename T>
struct InferMetaTypeTag {};

template <typename Fn, Fn fn>
struct InferMetaFnImpl;

template <typename Return, typename... Args, Return (*infer_meta_fn)(Args...)>
struct InferMetaFnImpl<Return (*)(Args...), infer_meta_fn> {
  static void Call(InferMetaContext* ctx) {
    InferMetaFnCallHelper<Args...,
                          InferMetaTypeTag<int>>::template Call<0, 0, 0>(ctx);
  }

 private:
  template <typename... RemainingArgs>
  struct InferMetaFnCallHelper;

  template <typename... Tail>
  struct InferMetaFnCallHelper<const MetaTensor&, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
      static_assert(attr_idx == 0,
                    "InferMeta's Input should appear before Attributes.");
      static_assert(out_idx == 0,
                    "InferMeta's Input should appear before Outputs.");
      const std::pair<int, int> range = ctx->InputRangeAt(in_idx);
      const MetaTensor& arg = ctx->InputAt(range.first);
      InferMetaFnCallHelper<
          Tail...>::template Call<in_idx + 1, attr_idx, out_idx>(ctx,
                                                                 pargs...,
                                                                 arg);
    }
  };

141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
  template <typename... Tail>
  struct InferMetaFnCallHelper<paddle::optional<const MetaTensor&>, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
      static_assert(attr_idx == 0,
                    "InferMeta's Input should appear before Attributes.");
      static_assert(out_idx == 0,
                    "InferMeta's Input should appear before Outputs.");
      const std::pair<int, int> range = ctx->InputRangeAt(in_idx);
      auto arg = ctx->OptionalInputAt(range.first);

      InferMetaFnCallHelper<
          Tail...>::template Call<in_idx + 1, attr_idx, out_idx>(ctx,
                                                                 pargs...,
                                                                 arg);
    }
  };

159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
  template <typename... Tail>
  struct InferMetaFnCallHelper<const std::vector<MetaTensor>&, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
      static_assert(attr_idx == 0,
                    "InferMeta's Input should appear before Attributes.");
      static_assert(out_idx == 0,
                    "InferMeta's Input should appear before Outputs.");
      const std::pair<int, int> range = ctx->InputRangeAt(in_idx);
      std::vector<MetaTensor> arg =
          ctx->InputsBetween(range.first, range.second);
      InferMetaFnCallHelper<
          Tail...>::template Call<in_idx + 1, attr_idx, out_idx>(ctx,
                                                                 pargs...,
                                                                 arg);
    }
  };

177
  // TODO(chenweihang): support other attr type later
178 179 180
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(bool);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(int);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(int64_t);
181
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(float);
182 183
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::string&);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::vector<bool>&);
184 185 186
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::vector<int>&);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(
      const std::vector<int64_t>&);
187 188 189 190
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::vector<float>&);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::vector<double>&);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(
      const std::vector<std::string>&);
191
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataType);
192
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(Backend);
193 194 195 196
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataLayout);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const Scalar&);
  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const ScalarArray&);

197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
  // TODO(chenweihang): support vector<MetaTensor> input later

  template <typename... Tail>
  struct InferMetaFnCallHelper<MetaTensor*, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
      const std::pair<int, int> range = ctx->OutputRangeAt(out_idx);
      MetaTensor* arg = ctx->MutableOutputAt(range.first);
      InferMetaFnCallHelper<
          Tail...>::template Call<in_idx, attr_idx, out_idx + 1>(ctx,
                                                                 pargs...,
                                                                 arg);
    }
  };

212 213 214 215 216 217 218 219 220 221 222 223 224 225
  template <typename... Tail>
  struct InferMetaFnCallHelper<std::vector<MetaTensor>*, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
      const std::pair<int, int> range = ctx->OutputRangeAt(out_idx);
      std::vector<MetaTensor> tmp =
          ctx->MutableOutputBetween(range.first, range.second);
      std::vector<MetaTensor>* arg = &tmp;
      InferMetaFnCallHelper<
          Tail...>::template Call<in_idx, attr_idx, out_idx + 1>(ctx,
                                                                 pargs...,
                                                                 arg);
    }
  };
226 227 228 229 230

  template <typename... Tail>
  struct InferMetaFnCallHelper<MetaConfig, Tail...> {
    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
231
      MetaConfig arg = ctx->GetMetaConfig();
232 233 234 235 236 237 238 239 240 241 242 243 244 245 246
      InferMetaFnCallHelper<Tail...>::template Call<in_idx, attr_idx, out_idx>(
          ctx, pargs..., arg);
    }
  };

  /* End case */
  template <typename T>
  struct InferMetaFnCallHelper<InferMetaTypeTag<T>> {
    template <int in_idx, int attr_idx, int out_idx>
    static void Call(InferMetaContext* ctx, Args&... args) {
      return infer_meta_fn(args...);
    }
  };
};

C
Chen Weihang 已提交
247
class MetaFnFactory {
248
 public:
C
Chen Weihang 已提交
249
  static MetaFnFactory& Instance();
250 251 252 253 254 255 256 257 258

  bool Contains(const std::string& kernel_name_prefix) const {
    return meta_fn_map_.count(kernel_name_prefix) > 0;
  }

  void Insert(std::string kernel_name_prefix, InferMetaFn infer_meta_fn) {
    PADDLE_ENFORCE_NE(
        Contains(kernel_name_prefix),
        true,
259
        phi::errors::AlreadyExists(
260 261 262 263 264 265 266 267 268 269 270
            "`%s`'s Series Kernel's InferMetaFn has been registered.",
            kernel_name_prefix));
    meta_fn_map_.insert(
        {std::move(kernel_name_prefix), std::move(infer_meta_fn)});
  }

  const InferMetaFn& Get(const std::string& kernel_name_prefix) const {
    auto it = meta_fn_map_.find(kernel_name_prefix);
    PADDLE_ENFORCE_NE(
        it,
        meta_fn_map_.end(),
271
        phi::errors::NotFound(
272 273 274 275 276 277
            "`%s`'s Series Kernel's InferMetaFn is not registered.",
            kernel_name_prefix));
    return it->second;
  }

 private:
C
Chen Weihang 已提交
278
  MetaFnFactory() = default;
279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295

  /**
   * [ Why use kernel name prefix? ]
   *
   * one op -> a matrix of kernels
   *
   * such as, scale op, it may correspond to the following kernels:
   *
   * - scale, scale_sr, scale_dnnl
   * - scale_raw, scale_raw_sr, scale_raw_dnnl
   *
   * All the kernels in each row correspond to the same infershape function,
   * the number of kernel arguments in the same row is the same, and only
   * the tensor types in the arguments are different.
   */
  paddle::flat_hash_map<std::string, InferMetaFn> meta_fn_map_;

C
Chen Weihang 已提交
296
  DISABLE_COPY_AND_ASSIGN(MetaFnFactory);
297 298 299 300 301
};

struct InferMetaFnRegistrar {
  InferMetaFnRegistrar(const char* kernel_name_prefix,
                       InferMetaFn infer_meta_fn) {
C
Chen Weihang 已提交
302 303
    MetaFnFactory::Instance().Insert(kernel_name_prefix,
                                     std::move(infer_meta_fn));
304 305 306
  }
};

307
#define PD_REGISTER_INFER_META_FN(kernel_name_prefix, variadic_infer_meta_fn) \
308
  PT_STATIC_ASSERT_GLOBAL_NAMESPACE(                                          \
309 310
      PD_REGISTER_infer_meta_fn_ns_check_##kernel_name_prefix,                \
      "PD_REGISTER_INFER_META_FN must be called in global namespace.");       \
311
  static const ::phi::InferMetaFnRegistrar                                    \
312
      __registrar_arg_map_fn_for_##kernel_name_prefix(                        \
313
          #kernel_name_prefix, PT_INFER_META(variadic_infer_meta_fn))
314

315
}  // namespace phi