custom_operator.cc 27.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
/* 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. */

#include "paddle/fluid/framework/custom_operator.h"

#include <algorithm>
#include <functional>
#include <iostream>
#include <map>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "paddle/fluid/extension/include/tensor.h"
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/c/c_api.h"
#include "paddle/fluid/framework/custom_tensor_utils.h"
#include "paddle/fluid/framework/framework.pb.h"
#include "paddle/fluid/framework/op_meta_info_helper.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/platform/dynload/dynamic_loader.h"
#include "paddle/fluid/string/string_helper.h"

namespace paddle {
namespace framework {

namespace detail {

// dynamic lib load func
template <typename T>
static T* DynLoad(void* handle, std::string name) {
  T* func = reinterpret_cast<T*>(dlsym(handle, name.c_str()));
#if !defined(_WIN32)
  auto errorno = dlerror();
#else
  auto errorno = GetLastError();
#endif  // !_WIN32
  PADDLE_ENFORCE_NOT_NULL(
      func, platform::errors::NotFound(
                "Failed to load dynamic operator library, error message(%s).",
                errorno));
  return func;
}

inline bool IsGradVar(const std::string& var_name) {
  std::string suffix = kGradVarSuffix;
  return var_name.rfind(suffix) != std::string::npos;
}

inline std::string NoGrad(const std::string& var_name) {
  std::string suffix = kGradVarSuffix;
  return var_name.substr(0, var_name.size() - kGradVarSuffixSize);
}

inline bool IsMemberOf(const std::vector<std::string>& vec,
                       const std::string& name) {
  return std::find(vec.cbegin(), vec.cend(), name) != vec.cend();
}

76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
std::vector<std::string> ParseAttrStr(const std::string& attr) {
  auto split_pos = attr.find_first_of(":");
  PADDLE_ENFORCE_NE(split_pos, std::string::npos,
                    platform::errors::InvalidArgument(
                        "Invalid attribute string format. Attribute string "
                        "format is `<name>:<type>`."));

  std::vector<std::string> rlt;
  // 1. name
  rlt.emplace_back(string::trim_spaces(attr.substr(0, split_pos)));
  // 2. type
  rlt.emplace_back(string::trim_spaces(attr.substr(split_pos + 1)));

  VLOG(1) << "attr name: " << rlt[0] << ", attr type str: " << rlt[1];

  return rlt;
}

94 95 96 97 98 99 100 101
}  // namespace detail

////////////////// Kernel Define ////////////////////

// custom op kernel call function define
static void RunKernelFunc(const framework::ExecutionContext& ctx,
                          const paddle::KernelFunc& func,
                          const std::vector<std::string>& inputs,
102 103
                          const std::vector<std::string>& outputs,
                          const std::vector<std::string>& attrs) {
104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119
  VLOG(1) << "Custom Operator: Start run KernelFunc.";
  std::vector<paddle::Tensor> custom_ins;
  for (auto& in_name : inputs) {
    VLOG(1) << "Custom Operator: input name - " << in_name;
    auto* x = ctx.Input<Tensor>(in_name);
    PADDLE_ENFORCE_NOT_NULL(x, platform::errors::NotFound(
                                   "Input tensor (%s) is nullptr.", in_name));
    PADDLE_ENFORCE_EQ(x->IsInitialized(), true,
                      platform::errors::InvalidArgument(
                          "Input tensor (%s) is not initialized."));
    auto custom_in = paddle::Tensor(
        CustomTensorUtils::ConvertInnerPlaceToEnumPlace(x->place()));
    CustomTensorUtils::ShareDataFrom(static_cast<const void*>(x), custom_in);
    custom_ins.emplace_back(custom_in);
  }

120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
  std::vector<boost::any> custom_attrs;
  for (auto& attr_str : attrs) {
    auto attr_name_and_type = detail::ParseAttrStr(attr_str);
    auto attr_name = attr_name_and_type[0];
    auto attr_type_str = attr_name_and_type[1];
    if (attr_type_str == "bool") {
      custom_attrs.emplace_back(ctx.Attr<bool>(attr_name));
    } else if (attr_type_str == "int") {
      custom_attrs.emplace_back(ctx.Attr<int>(attr_name));
    } else if (attr_type_str == "float") {
      custom_attrs.emplace_back(ctx.Attr<float>(attr_name));
    } else if (attr_type_str == "int64_t") {
      custom_attrs.emplace_back(ctx.Attr<int64_t>(attr_name));
    } else if (attr_type_str == "std::string") {
      custom_attrs.emplace_back(ctx.Attr<std::string>(attr_name));
    } else if (attr_type_str == "std::vector<int>") {
      custom_attrs.emplace_back(ctx.Attr<std::vector<int>>(attr_name));
    } else if (attr_type_str == "std::vector<float>") {
      custom_attrs.emplace_back(ctx.Attr<std::vector<float>>(attr_name));
    } else if (attr_type_str == "std::vector<int64_t>") {
      custom_attrs.emplace_back(ctx.Attr<std::vector<int64_t>>(attr_name));
    } else if (attr_type_str == "std::vector<std::string>") {
      custom_attrs.emplace_back(ctx.Attr<std::vector<std::string>>(attr_name));
    } else {
      PADDLE_THROW(platform::errors::Unimplemented(
          "Unsupported `%s` type value as custom attribute now. "
          "Supported data types include `bool`, `int`, `float`, "
          "`int64_t`, `std::string`, `std::vector<int>`, "
          "`std::vector<float>`, `std::vector<int64_t>, "
          "`std::vector<std::string>`, Please check whether "
          "the attribute data type and data type string are matched.",
          attr_type_str));
    }
  }
154 155

  VLOG(1) << "Run ComputeFunc.";
156 157
  try {
    auto outs = func(custom_ins, custom_attrs);
158

159 160 161 162 163 164 165 166 167 168 169 170
    VLOG(1) << "Custom Operator: Share outputs into ExecutionContext.";
    for (size_t i = 0; i < outputs.size(); ++i) {
      auto* true_out = ctx.Output<Tensor>(outputs[i]);
      CustomTensorUtils::ShareDataTo(outs.at(i), true_out);
    }
  } catch (platform::EnforceNotMet& exception) {
    throw std::move(exception);
  } catch (std::exception& ex) {
    PADDLE_THROW(platform::errors::External("%s", ex.what()));
  } catch (...) {
    PADDLE_THROW(platform::errors::Fatal(
        "Custom operator raises an unknown exception in rumtime."));
171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
  }
}

//////////////////// Operator Define /////////////////

class CustomOperator : public OperatorWithKernel {
 public:
  using OperatorWithKernel::OperatorWithKernel;

  // Dummy infershape
  // Because it is a pure virtual function, it must be implemented
  void InferShape(framework::InferShapeContext* ctx) const override {
    VLOG(1) << "Custom Operator: Dummy infer shape of custom operator.";
  }

  /**
   * NOTE: [Skip the Kernel Selection]
   * Custom Op only registers one Op kernel on each device, so that the
   * data type selection and promotion that depends on GetExpectedKernelType,
   * as well as the adaptation of various other special situations,
   * need users to implement, to avoid users needs to implement
   * GetExpectedKernelType function when expanding other cases.
   * The RAW type is used here as the data type, indicating that
   * it can only be determined at runtime.
   */
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const {
    return framework::OpKernelType(proto::VarType::RAW, ctx.GetPlace());
  }

  /**
   * NOTE: [Skip Input Variable Cast for DataType]
   * Because the kernel data type is RAW, we should skip the cast for
   * data type difference when PrepareData.
   */
  framework::OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const Tensor& tensor,
      const OpKernelType& expected_kernel_type) {
    return OpKernelType(expected_kernel_type.data_type_,
                        expected_kernel_type.place_, tensor.layout());
  }
};

class CustomOpMaker : public OpProtoAndCheckerMaker {
 public:
  explicit CustomOpMaker(const std::vector<std::string>& inputs,
                         const std::vector<std::string>& outputs,
                         const std::vector<std::string>& attrs)
      : inputs_(inputs), outputs_(outputs), attrs_(attrs) {}

  void Make() override {
    for (auto& in_name : inputs_) {
      AddInput(in_name, "The input " + in_name + "of Custom operator.");
    }
    for (auto& out_name : outputs_) {
      AddOutput(out_name, "The output " + out_name + "of Custom Operator.");
    }
228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272
    for (auto& attr : attrs_) {
      auto attr_name_and_type = detail::ParseAttrStr(attr);
      auto attr_name = attr_name_and_type[0];
      auto attr_type_str = attr_name_and_type[1];
      if (attr_type_str == "bool") {
        AddAttr<bool>(attr_name, "custom operator bool attribute.")
            .SetDefault(false);
      } else if (attr_type_str == "int") {
        AddAttr<int>(attr_name, "custom operator int attribute.").SetDefault(1);
      } else if (attr_type_str == "float") {
        AddAttr<float>(attr_name, "custom operator float attribute.")
            .SetDefault(1.0f);
      } else if (attr_type_str == "int64_t") {
        AddAttr<int64_t>(attr_name, "custom operator int64_t attribute.")
            .SetDefault(1);
      } else if (attr_type_str == "std::string") {
        AddAttr<std::string>(attr_name, "custom operator int attribute.")
            .SetDefault("");
      } else if (attr_type_str == "std::vector<int>") {
        AddAttr<std::vector<int>>(attr_name,
                                  "custom operator std::vector<int> attribute.")
            .SetDefault({});
      } else if (attr_type_str == "std::vector<float>") {
        AddAttr<std::vector<float>>(
            attr_name, "custom operator std::vector<float> attribute.")
            .SetDefault({});
      } else if (attr_type_str == "std::vector<int64_t>") {
        AddAttr<std::vector<int64_t>>(
            attr_name, "custom operator std::vector<int64_t> attribute.")
            .SetDefault({});
      } else if (attr_type_str == "std::vector<std::string>") {
        AddAttr<std::vector<std::string>>(
            attr_name, "custom operator std::vector<std::string> attribute.")
            .SetDefault({});
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported `%s` type value as custom attribute now. "
            "Supported data types include `bool`, `int`, `float`, "
            "`int64_t`, `std::string`, `std::vector<int>`, "
            "`std::vector<float>`, `std::vector<int64_t>, "
            "`std::vector<std::string>`, Please check whether "
            "the attribute data type and data type string are matched.",
            attr_type_str));
      }
    }
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334
    AddComment(R"DOC(
Custom Operator.

According to the Tensor operation function implemented by the user 
independently of the framework, it is encapsulated into a framework 
operator to adapt to various execution scenarios such as dynamic graph, 
mode static graph mode, and inference mode.

)DOC");
  }

 private:
  std::vector<std::string> inputs_;
  std::vector<std::string> outputs_;
  std::vector<std::string> attrs_;
};

template <typename T>
class CustomGradOpMaker;

template <>
class CustomGradOpMaker<OpDesc> : public SingleGradOpMaker<OpDesc> {
 public:
  explicit CustomGradOpMaker(
      const OpDesc& fwd_op, const std::unordered_set<std::string>& no_grad_set,
      std::unordered_map<std::string, std::string>* grad_to_var,
      const std::vector<BlockDesc*>& grad_block, const std::string& name,
      const std::vector<std::string>& inputs,
      const std::vector<std::string>& outputs)
      : SingleGradOpMaker<OpDesc>(fwd_op, no_grad_set, grad_to_var, grad_block),
        name_(name),
        inputs_(inputs),
        outputs_(outputs) {}

 protected:
  void Apply(GradOpPtr<OpDesc> grad_op) const override {
    grad_op->SetType(name_);

    auto fwd_op_inputs = this->InputNames();
    auto fwd_op_outputs = this->OutputNames();

    for (auto& in_name : inputs_) {
      VLOG(1) << "Custom Operator: GradOpDescMaker - input: " << in_name;
      if (!detail::IsGradVar(in_name)) {
        if (detail::IsMemberOf(fwd_op_inputs, in_name)) {
          grad_op->SetInput(in_name, this->Input(in_name));
        } else if (detail::IsMemberOf(fwd_op_outputs, in_name)) {
          grad_op->SetInput(in_name, this->Output(in_name));
        } else {
          PADDLE_THROW(platform::errors::InvalidArgument(
              "The input tensor name `%s` is invalid, expected it is the input "
              "or output of forward operator.",
              in_name));
        }
      } else {
        grad_op->SetInput(in_name, this->OutputGrad(detail::NoGrad(in_name)));
      }
    }
    for (auto& out_name : outputs_) {
      VLOG(1) << "Custom Operator: GradOpDescMaker - output: " << out_name;
      grad_op->SetOutput(out_name, this->InputGrad(detail::NoGrad(out_name)));
    }
335
    grad_op->SetAttrMap(this->Attrs());
336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394
  }

 private:
  std::string name_;
  std::vector<std::string> inputs_;
  std::vector<std::string> outputs_;
};

template <>
class CustomGradOpMaker<imperative::OpBase>
    : public SingleGradOpMaker<imperative::OpBase> {
 public:
  explicit CustomGradOpMaker(
      const std::string& type,
      const imperative::NameVarBaseMap& var_base_map_in,
      const imperative::NameVarBaseMap& var_base_map_out,
      const AttributeMap& attrs,
      const std::map<std::string, std::string>& inplace_map,
      const std::string& name, const std::vector<std::string>& inputs,
      const std::vector<std::string>& outputs)
      : SingleGradOpMaker<imperative::OpBase>(
            type, var_base_map_in, var_base_map_out, attrs, inplace_map),
        name_(name),
        inputs_(inputs),
        outputs_(outputs) {}

 protected:
  // TODO(chenweihang): The code is duplicated with the previous one, because
  // ere OpMaker's Input, Output and other methods are protected. Putting the
  // function implementation outside the class will cause the method to be
  // uncallable,
  // so it is still implemented in the class for the time being.
  void Apply(GradOpPtr<imperative::OpBase> grad_op) const override {
    grad_op->SetType(name_);

    auto fwd_op_inputs = this->InputNames();
    auto fwd_op_outputs = this->OutputNames();

    for (auto& in_name : inputs_) {
      VLOG(1) << "Custom Operator: GradOpBaseMaker - input: " << in_name;
      if (!detail::IsGradVar(in_name)) {
        if (detail::IsMemberOf(fwd_op_inputs, in_name)) {
          grad_op->SetInput(in_name, this->Input(in_name));
        } else if (detail::IsMemberOf(fwd_op_outputs, in_name)) {
          grad_op->SetInput(in_name, this->Output(in_name));
        } else {
          PADDLE_THROW(platform::errors::InvalidArgument(
              "The input tensor name `%s` is invalid, expected it is the input "
              "or output of forward operator.",
              in_name));
        }
      } else {
        grad_op->SetInput(in_name, this->OutputGrad(detail::NoGrad(in_name)));
      }
    }
    for (auto& out_name : outputs_) {
      VLOG(1) << "Custom Operator: GradOpBaseMaker - output: " << out_name;
      grad_op->SetOutput(out_name, this->InputGrad(detail::NoGrad(out_name)));
    }
395
    grad_op->SetAttrMap(this->Attrs());
396 397 398 399 400 401 402 403 404 405 406 407 408 409 410
  }

 private:
  std::string name_;
  std::vector<std::string> inputs_;
  std::vector<std::string> outputs_;
};

//////////// Operator and Kernel Register //////////////

void RegisterOperatorKernelWithPlace(const std::string& name,
                                     const paddle::KernelFunc& kernel_func,
                                     const proto::VarType::Type type,
                                     const PlaceType& place,
                                     const std::vector<std::string>& inputs,
411 412
                                     const std::vector<std::string>& outputs,
                                     const std::vector<std::string>& attrs) {
413 414 415 416
  OpKernelType key(type,
                   CustomTensorUtils::ConvertEnumPlaceToInnerPlace(place));
  VLOG(1) << "Custom Operator: op kernel key: " << key;
  OperatorWithKernel::AllOpKernels()[name][key] =
417 418
      [kernel_func, inputs, outputs,
       attrs](const framework::ExecutionContext& ctx) {
419
        VLOG(1) << "Custom Operator: run custom kernel func in lambda.";
420
        RunKernelFunc(ctx, kernel_func, inputs, outputs, attrs);
421 422 423 424 425 426
      };
}

void RegisterOperatorKernel(const std::string& name,
                            const paddle::KernelFunc& kernel_func,
                            const std::vector<std::string>& inputs,
427 428
                            const std::vector<std::string>& outputs,
                            const std::vector<std::string>& attrs) {
429 430 431 432 433 434 435
  VLOG(1) << "Custom Operator: op name in kernel: " << name;
  // NOTE [ Dummy Op Kernel Key ]
  // TODO(chenweihang): Because execute engine need get device context based
  // op_kernel_key.place_, so we should register kernel for each
  // device. But this is not entirely correct, if user only give a cpu kernel,
  // but call api in gpu device, it will cause error.
  RegisterOperatorKernelWithPlace(name, kernel_func, proto::VarType::RAW,
436 437
                                  PlaceType::kCPU, inputs, outputs, attrs);
#ifdef PADDLE_WITH_CUDA
438
  RegisterOperatorKernelWithPlace(name, kernel_func, proto::VarType::RAW,
439 440
                                  PlaceType::kGPU, inputs, outputs, attrs);
#endif
441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462
}

void RegisterOperatorWithMetaInfo(
    const std::vector<OpMetaInfo>& op_meta_infos) {
  /* Op register */
  OpInfo info;

  auto& base_op_meta = op_meta_infos.front();

  auto op_name = OpMetaInfoHelper::GetOpName(base_op_meta);
  auto& op_inputs = OpMetaInfoHelper::GetInputs(base_op_meta);
  auto& op_outputs = OpMetaInfoHelper::GetOutputs(base_op_meta);
  auto& op_attrs = OpMetaInfoHelper::GetAttrs(base_op_meta);
  auto& kernel_fn = OpMetaInfoHelper::GetKernelFn(base_op_meta);
  auto& infer_shape_func = OpMetaInfoHelper::GetInferShapeFn(base_op_meta);
  auto& infer_dtype_func = OpMetaInfoHelper::GetInferDtypeFn(base_op_meta);

  VLOG(1) << "Custom Operator: forward, op name: " << op_name;
  VLOG(1) << "Custom Operator: forward, op inputs: "
          << string::join_strings(op_inputs, ',');
  VLOG(1) << "Custom Operator: forward, op outputs: "
          << string::join_strings(op_outputs, ',');
463 464
  VLOG(1) << "Custom Operator: forward, op attrs: "
          << string::join_strings(op_attrs, ',');
465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486

  // Op
  info.creator_ = [](const std::string& op_name, const VariableNameMap& inputs,
                     const VariableNameMap& outputs,
                     const AttributeMap& attrs) {
    return new CustomOperator(op_name, inputs, outputs, attrs);
  };

  // OpMaker
  info.proto_ = new proto::OpProto;
  info.proto_->set_type(op_name);

  info.checker_ = new OpAttrChecker();
  CustomOpMaker custom_maker(op_inputs, op_outputs, op_attrs);
  custom_maker(info.proto_, info.checker_);
  PADDLE_ENFORCE_EQ(
      info.proto_->IsInitialized(), true,
      platform::errors::PreconditionNotMet(
          "Fail to initialize %s's OpProto, because %s is not initialized.",
          op_name, info.proto_->InitializationErrorString()));

  // InferShape
487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524
  if (infer_shape_func == nullptr) {
    // use default InferShape
    info.infer_shape_ = [op_inputs, op_outputs](InferShapeContext* ctx) {
      PADDLE_ENFORCE_EQ(
          op_inputs.size(), 1UL,
          platform::errors::Unavailable(
              "Your custom operator contains multiple inputs. "
              "We only allow a custom operator that contains only one input "
              "and "
              "only one output without setting the InferShapeFn. At this time, "
              "the input shape will be directly set to the output shape.\n"
              "Please set the InferShapeFn of custom "
              "operator by .SetInferShapeFn(PD_INFER_SHAPE(...))"));
      PADDLE_ENFORCE_EQ(
          op_outputs.size(), 1UL,
          platform::errors::Unavailable(
              "Your custom operator contains multiple outputs. "
              "We only allow a custom operator that contains only one input "
              "and "
              "only one output without setting the InferShapeFn. At this time, "
              "the input shape will be directly set to the output shape.\n"
              "Please set the InferShapeFn of custom "
              "operator by .SetInferShapeFn(PD_INFER_SHAPE(...))"));

      VLOG(1) << "Custom Operator: Default InferShape - share ddim.";
      ctx->ShareDim(op_inputs[0], op_outputs[0]);
    };
  } else {
    info.infer_shape_ = [op_inputs, op_outputs,
                         infer_shape_func](InferShapeContext* ctx) {
      std::vector<std::vector<int64_t>> input_shapes;

      VLOG(1) << "Custom Operator: InferShape - get input ddim.";
      for (auto& in_name : op_inputs) {
        OP_INOUT_CHECK(ctx->HasInput(in_name), "Input", in_name, "Custom");
        auto ddim = ctx->GetInputDim(in_name);
        input_shapes.emplace_back(framework::vectorize(ddim));
      }
525

526 527
      VLOG(1) << "Custom Operator: InferShape - calc output ddim.";
      auto output_shapes = infer_shape_func(input_shapes);
528

529 530 531 532 533 534 535
      VLOG(1) << "Custom Operator: InferShape - set output ddim.";
      for (size_t i = 0; i < op_outputs.size(); ++i) {
        ctx->SetOutputDim(op_outputs[i],
                          framework::make_ddim(output_shapes[i]));
      }
    };
  }
536 537

  // Infer Dtype
538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576
  if (infer_dtype_func == nullptr) {
    // use defalut InferDtype
    info.infer_var_type_ = [op_inputs, op_outputs](InferVarTypeContext* ctx) {
      PADDLE_ENFORCE_EQ(
          op_inputs.size(), 1UL,
          platform::errors::Unavailable(
              "Your custom operator contains multiple inputs. "
              "We only allow a custom operator that contains only one input "
              "and "
              "only one output without setting the InferDtypeFn. At this time, "
              "the input dtype will be directly set to the output dtype.\n"
              "Please set the InferDtypeFn of custom "
              "operator by .SetInferDtypeFn(PD_INFER_DTYPE(...))"));
      PADDLE_ENFORCE_EQ(
          op_outputs.size(), 1UL,
          platform::errors::Unavailable(
              "Your custom operator contains multiple outputs. "
              "We only allow a custom operator that contains only one input "
              "and "
              "only one output without setting the InferDtypeFn. At this time, "
              "the input dtype will be directly set to the output dtype.\n"
              "Please set the InferDtypeFn of custom "
              "operator by .SetInferDtypeFn(PD_INFER_DTYPE(...))"));

      VLOG(1) << "Custom Operator: InferDtype - share dtype.";
      auto dtype = ctx->GetInputDataType(op_inputs[0]);
      ctx->SetOutputDataType(op_outputs[0], dtype);
    };
  } else {
    info.infer_var_type_ = [op_inputs, op_outputs,
                            infer_dtype_func](InferVarTypeContext* ctx) {
      std::vector<DataType> input_dtypes;

      VLOG(1) << "Custom Operator: InferDtype - get input dtype.";
      for (auto& in_name : op_inputs) {
        auto dtype = ctx->GetInputDataType(in_name);
        input_dtypes.emplace_back(
            CustomTensorUtils::ConvertInnerDTypeToEnumDType(dtype));
      }
577

578 579
      VLOG(1) << "Custom Operator: InferDtype - infer output dtype.";
      auto output_dtypes = infer_dtype_func(input_dtypes);
580

581 582 583 584 585 586 587 588
      VLOG(1) << "Custom Operator: InferDtype - set output dtype.";
      for (size_t i = 0; i < op_outputs.size(); ++i) {
        ctx->SetOutputDataType(
            op_outputs[i],
            CustomTensorUtils::ConvertEnumDTypeToInnerDType(output_dtypes[i]));
      }
    };
  }
589 590

  // Kernel func
591
  RegisterOperatorKernel(op_name, kernel_fn, op_inputs, op_outputs, op_attrs);
592 593 594 595 596 597 598 599 600

  // If grad op or double grad op exists
  std::string cur_op_name = op_name;
  for (size_t i = 1; i < op_meta_infos.size(); ++i) {
    auto& cur_grad_op = op_meta_infos[i];

    auto& grad_op_name = OpMetaInfoHelper::GetOpName(cur_grad_op);
    auto& grad_op_inputs = OpMetaInfoHelper::GetInputs(cur_grad_op);
    auto& grad_op_outputs = OpMetaInfoHelper::GetOutputs(cur_grad_op);
601
    auto& grad_op_attrs = OpMetaInfoHelper::GetAttrs(cur_grad_op);
602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654
    auto& grad_kernel_fn = OpMetaInfoHelper::GetKernelFn(cur_grad_op);

    VLOG(1) << "Custom Operator: backward, op name: " << grad_op_name;
    VLOG(1) << "Custom Operator: backward, op inputs: "
            << string::join_strings(grad_op_inputs, ',');
    VLOG(1) << "Custom Operator: backward, op outputs: "
            << string::join_strings(grad_op_outputs, ',');

    // GradOpDescMaker
    info.grad_op_maker_ = [grad_op_name, grad_op_inputs, grad_op_outputs](
        const OpDesc& fwd_op,
        const std::unordered_set<std::string>& no_grad_set,
        std::unordered_map<std::string, std::string>* grad_to_var,
        const std::vector<BlockDesc*>& grad_block) {
      CustomGradOpMaker<paddle::framework::OpDesc> maker(
          fwd_op, no_grad_set, grad_to_var, grad_block, grad_op_name,
          grad_op_inputs, grad_op_outputs);
      return maker();
    };

    // GradOpBaseMaker
    info.dygraph_grad_op_maker_ = [grad_op_name, grad_op_inputs,
                                   grad_op_outputs](
        const std::string& type,
        const imperative::NameVarBaseMap& var_base_map_in,
        const imperative::NameVarBaseMap& var_base_map_out,
        const framework::AttributeMap& attrs,
        const std::map<std::string, std::string>& inplace_map) {
      CustomGradOpMaker<paddle::imperative::OpBase> maker(
          type, var_base_map_in, var_base_map_out, attrs, inplace_map,
          grad_op_name, grad_op_inputs, grad_op_outputs);
      return maker();
    };

    /* Grad op register */
    OpInfo grad_info;

    // Grad Op
    grad_info.creator_ = [](
        const std::string& type, const VariableNameMap& inputs,
        const VariableNameMap& outputs, const AttributeMap& attrs) {
      return new CustomOperator(type, inputs, outputs, attrs);
    };

    // Grad InferShape (gradient's shape is same with forward input default)
    grad_info.infer_shape_ = [grad_op_outputs](InferShapeContext* ctx) {
      for (auto& out_name : grad_op_outputs) {
        ctx->ShareDim(detail::NoGrad(out_name), out_name);
      }
    };

    // Kernel func
    RegisterOperatorKernel(grad_op_name, grad_kernel_fn, grad_op_inputs,
655
                           grad_op_outputs, grad_op_attrs);
656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697

    // update current info
    OpInfoMap::Instance().Insert(cur_op_name, info);
    cur_op_name = grad_op_name;
    info = grad_info;
  }
  // insert last info
  OpInfoMap::Instance().Insert(cur_op_name, info);
}

void RegisterOperatorWithMetaInfoMap(
    const paddle::OpMetaInfoMap& op_meta_info_map) {
  auto& meta_info_map = op_meta_info_map.GetMap();

  PADDLE_ENFORCE_EQ(meta_info_map.empty(), false,
                    platform::errors::PreconditionNotMet(
                        "No custom operator that needs to be registered."));
  VLOG(1) << "Custom Operator: size of op meta info map - "
          << meta_info_map.size();
  // pair: {op_type, OpMetaInfo}
  for (auto& pair : meta_info_map) {
    VLOG(1) << "Custom Operator: pair first -> op name: " << pair.first;
    RegisterOperatorWithMetaInfo(pair.second);
  }
}

////////////////////// User APIs ///////////////////////

// load op api
void LoadOpMetaInfoAndRegisterOp(const std::string& dso_name) {
  void* handle = paddle::platform::dynload::GetOpDsoHandle(dso_name);

  typedef OpMetaInfoMap& get_op_meta_info_map_t();
  auto* get_op_meta_info_map =
      detail::DynLoad<get_op_meta_info_map_t>(handle, "PD_GetOpMetaInfoMap");
  auto& op_meta_info_map = get_op_meta_info_map();

  RegisterOperatorWithMetaInfoMap(op_meta_info_map);
}

}  // namespace framework
}  // namespace paddle