op_function_generator.cc 14.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// Copyright (c) 2019 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.

15
#include <algorithm>
16 17 18 19 20 21 22 23 24 25 26
#include <fstream>
#include <iostream>
#include <string>

#include "paddle/fluid/framework/op_info.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/variable.h"
#include "paddle/fluid/pybind/pybind.h"
#include "paddle/fluid/string/string_helper.h"

L
Leo Chen 已提交
27 28 29 30 31 32 33 34
// NOTE(zhiqiu): Commonly, the inputs in auto-generated OP function are
// determined by the OP`s proto automatically, i.e., all the inputs registered
// in OpMaker.
// However, some OPs have dispensable inputs, which means the input can
// be none for some conditions. It is discovered that most dispensable inputs
// is not used in imperative mode, so we drop those inputs when generating OP
// functions. While, for very few OPs, the dispensable inputs are used, we
// need to manually specify them in this map.
35 36
std::map<std::string, std::set<std::string>> op_ins_map = {
    {"layer_norm", {"X", "Scale", "Bias"}},
C
ceci3 已提交
37
    {"instance_norm", {"X", "Scale", "Bias"}},
38 39 40
    {"gru_unit", {"Input", "HiddenPrev", "Weight", "Bias"}},
    {"label_smooth", {"X", "PriorDist"}},
    {"assign", {"X"}},
L
Leo Chen 已提交
41 42
    {"fake_quantize_dequantize_moving_average_abs_max",
     {"X", "InScale", "InAccum", "InState"}},
43
    {"nll_loss", {"X", "Label", "Weight"}},
44
};
L
Leo Chen 已提交
45 46 47 48 49 50 51 52 53 54 55 56

// NOTE(zhiqiu): Like op_ins_map.
// Commonly, the outputs in auto-generated OP function are determined by the
// OP`s proto automatically, i.e., all the outputs registered in OpMaker.
// However, some OPs have dispensable outputs, which means the output can
// be none for some conditions. It is discovered that most dispensable outputs
// is not used in imperative mode, so we drop those outputs when generating OP
// functions. While, for very few OPs, the dispensable outputs are used, we
// need to manually specify them in this map.
std::map<std::string, std::set<std::string>> op_outs_map = {
    {"fake_quantize_dequantize_moving_average_abs_max",
     {"Out", "OutScale", "OutAccum", "OutState"}},
57 58 59
    {"batch_norm",
     {"Y", "MeanOut", "VarianceOut", "SavedMean", "SavedVariance",
      "ReserveSpace"}},
L
Leo Chen 已提交
60 61 62 63 64 65 66 67 68 69 70 71 72 73
};

// NOTE(zhiqiu): Commonly, the outputs in auto-generated OP function are
// generated in C++ automatically.
// However, some OPs need to pass the outputs from Python instead of generating
// them in C++. There are mainly 2 reasons for that,
// (1) Optimizer OPs need to update the input param in-place, like sgd.
//     So they need to pass the output which is same as input param.
// (2) Very few python APIs has out in their arguments, like fill_constant.
//     So they need to pass the python output to C++.
//     Actually, this is not a good design, since it may break the SSA graph,
//     especially in declarative mode.
// For those OPs, we need to manually specify the outs need to pass in this map.
std::map<std::string, std::set<std::string>> op_passing_outs_map = {
74 75 76 77 78 79
    {"sgd", {"ParamOut"}},
    {"adam",
     {"ParamOut", "Moment1Out", "Moment2Out", "Beta1PowOut", "Beta2PowOut"}},
    {"momentum", {"ParamOut", "VelocityOut"}},
    {"batch_norm", {"MeanOut", "VarianceOut"}},
    {"accuracy", {"Correct", "Total"}},
80
    {"fill_constant", {"Out"}},
L
Leo Chen 已提交
81 82
    {"matmul", {"Out"}},
    {"fake_quantize_dequantize_moving_average_abs_max",
83
     {"Out", "OutScale", "OutAccum", "OutState"}},
84
    {"fake_quantize_dequantize_abs_max", {"Out", "OutScale"}},
85
    {"amp_check_finite_and_scale", {"Out", "FoundInfinite"}},
L
Leo Chen 已提交
86
};
87

88
// clang-format off
89 90
const char* OUT_INITIALIZER_TEMPLATE =
    R"({"%s", {std::shared_ptr<imperative::VarBase>(new imperative::VarBase(tracer->GenerateUniqueName()))}})";
91 92 93 94
const char* OUT_DUPLICABLE_INITIALIZER_TEMPLATE = R"({"%s", ConstructDuplicableOutput(%s)})";

const char* INPUT_INITIALIZER_TEMPLATE = R"({"%s", {%s}})";
const char* INPUT_LIST_INITIALIZER_TEMPLATE = R"({"%s", %s})";
L
Leo Chen 已提交
95 96 97 98 99

const char* INPUT_INITIALIZER_TEMPLATE_WITH_NULL = R"(	
    if (%s != nullptr) {	
      ins["%s"] = {%s};	
    }	
100
)";
L
Leo Chen 已提交
101 102 103 104 105 106 107

const char* INPUT_INITIALIZER_TEMPLATE_WITH_NULL_LIST = R"(	
    if (%s.size() != 0) {
      ins["%s"] = %s;	
    }	
)";

108 109
const char* OUTPUT_INITIALIZER_TEMPLATE_WITH_NULL = R"(
    outs["%s"] = {%s};
110 111
)";

112 113
const char* OUTPUT_INITIALIZER_TEMPLATE_WITH_NULL_LIST = R"(
    outs["%s"] = %s;
L
Leo Chen 已提交
114
)";
115 116 117 118
// if inputs is list, no need {}
const char* ARG_OUT_NUM = R"(%sNum)";
const char* ARG_OUT_NUM_TYPE = R"(size_t )";

119 120 121 122 123 124 125 126 127 128 129 130 131
const char* IN_VAR_TYPE = R"(py::handle)";
const char* IN_VAR_LIST_TYPE = R"(py::handle)";

const char* OUT_VAR_TYPE = R"(std::shared_ptr<imperative::VarBase>)";
const char* OUT_VAR_LIST_TYPE = R"(std::vector<std::shared_ptr<imperative::VarBase>>)";

const char* CAST_VAR_TEMPLATE = R"(
  auto %s = CastPyHandleToVarBase("%s", "%s", %d, %s);)";

const char* CAST_VAR_LIST_TEMPLATE = R"(
  auto %s = CastPyHandleToVarBaseList("%s", "%s", %d, %s);)";


132 133 134 135 136 137 138 139 140 141
const char* ARG_TEMPLATE = R"(const %s& %s)";

const char* RETURN_TUPLE_TYPE = R"(std::tuple<%s>)";
const char* RETURN_TYPE = R"(%s)";
const char* RETURN_TUPLE_TEMPLATE = R"(std::make_tuple(%s))";
const char* RETURN_LIST_TEMPLATE = R"(outs["%s"])";
const char* RETURN_TEMPLATE = R"(outs["%s"][0])";

const char* FUNCTION_ARGS = R"(%s, const py::args& args)";
const char* FUNCTION_ARGS_NO_INPUT = R"(const py::args& args)";
142 143

const char* OP_FUNCTION_TEMPLATE =
144
R"(
145
%s %s(%s)
146
{
147
  %s
148
  framework::AttributeMap attrs;
149
  ConstructAttrMapFromPyArgs("%s", %d, &attrs, args);
150 151 152 153 154 155 156 157 158
  {
    py::gil_scoped_release release;
    auto tracer = imperative::GetCurrentTracer();
    imperative::NameVarBaseMap outs = %s;
    imperative::NameVarBaseMap ins = %s;
    %s
    tracer->TraceOp("%s", ins, outs, attrs);
    return %s; 
  }   
159
})";
160

161
const char* PYBIND_ITEM_TEMPLATE = R"(  %s.def("%s", &%s);)";
162

163
// clang-format on
L
Leo Chen 已提交
164 165
static inline bool FindInsMap(const std::string& op_type,
                              const std::string& in_name) {
166 167 168
  return op_ins_map[op_type].count(in_name);
}

L
Leo Chen 已提交
169 170 171 172 173 174 175 176
static inline bool FindOutsMap(const std::string& op_type,
                               const std::string& out_name) {
  return op_outs_map[op_type].count(out_name);
}

static inline bool FindPassingOutsMap(const std::string& op_type,
                                      const std::string& out_name) {
  return op_passing_outs_map[op_type].count(out_name);
177
}
178

179 180 181 182
static inline std::string TempName(const std::string& name) {
  return name + '_';
}

183 184
static std::tuple<std::vector<std::string>, std::vector<std::string>>
GenerateOpFunctions(const std::string& module_name) {
185 186
  auto& op_info_map = paddle::framework::OpInfoMap::Instance().map();

187
  std::vector<std::string> op_function_list, bind_function_list;
188 189
  auto& all_kernels = paddle::framework::OperatorWithKernel::AllOpKernels();

190 191 192 193 194 195 196
  for (auto& pair : op_info_map) {
    auto& op_info = pair.second;
    auto op_proto = op_info.proto_;
    if (op_proto == nullptr) {
      continue;
    }
    auto& op_type = op_proto->type();
197 198 199 200 201 202 203 204 205
    // Skip ooerator which is not inherit form OperatorWithKernel, like while,
    // since only OperatorWithKernel can run in dygraph mode.
    if (!all_kernels.count(op_type)) {
      continue;
    }
    std::string input_args = "";
    std::string ins_initializer = "{";
    std::string ins_initializer_with_null = "";
    std::string py_arg = "";
206
    int arg_idx = 0;
207
    int input_args_num = 0;
208
    std::string ins_cast_str = "";
209 210 211
    for (auto& input : op_proto->inputs()) {
      auto& in_name = input.name();
      // skip those dispensable inputs, like ResidualData in conv2d
L
Leo Chen 已提交
212
      if (input.dispensable() && !FindInsMap(op_type, in_name)) {
213 214
        continue;
      }
215 216 217
      const auto in_type = input.duplicable() ? IN_VAR_LIST_TYPE : IN_VAR_TYPE;
      auto input_arg =
          paddle::string::Sprintf(ARG_TEMPLATE, in_type, TempName(in_name));
218 219
      input_args += input_arg;
      input_args += ",";
220
      input_args_num++;
221 222 223 224 225
      const auto in_cast_type =
          input.duplicable() ? CAST_VAR_LIST_TEMPLATE : CAST_VAR_TEMPLATE;
      ins_cast_str +=
          paddle::string::Sprintf(in_cast_type, in_name, op_type, in_name,
                                  arg_idx++, TempName(in_name));
226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249

      if (input.dispensable()) {
        const auto in_template = input.duplicable()
                                     ? INPUT_INITIALIZER_TEMPLATE_WITH_NULL_LIST
                                     : INPUT_INITIALIZER_TEMPLATE_WITH_NULL;
        ins_initializer_with_null +=
            paddle::string::Sprintf(in_template, in_name, in_name, in_name);
      } else {
        const auto in_template = input.duplicable()
                                     ? INPUT_LIST_INITIALIZER_TEMPLATE
                                     : INPUT_INITIALIZER_TEMPLATE;
        ins_initializer +=
            paddle::string::Sprintf(in_template, in_name, in_name);
        ins_initializer += ",";
      }
    }
    if (ins_initializer.back() == ',') {
      ins_initializer.pop_back();
    }
    ins_initializer += "}";

    if (input_args.back() == ',') {
      input_args.pop_back();
    }
250 251 252

    // Generate outs initializer
    std::string outs_initializer = "{";
L
Leo Chen 已提交
253
    std::string outs_initializer_with_null = "";
254 255
    std::string return_type = "";
    std::string return_str = "";
256

257
    int outs_num = 0;
258
    for (auto& output : op_proto->outputs()) {
L
Leo Chen 已提交
259 260 261
      auto& out_name = output.name();
      // skip those dispensable oututs
      if (output.dispensable() && !FindOutsMap(op_type, out_name)) {
262 263
        continue;
      }
264 265
      const auto out_type =
          output.duplicable() ? OUT_VAR_LIST_TYPE : OUT_VAR_TYPE;
266 267
      const auto return_template =
          output.duplicable() ? RETURN_LIST_TEMPLATE : RETURN_TEMPLATE;
L
Leo Chen 已提交
268
      if (FindPassingOutsMap(op_type, out_name)) {
269 270 271 272 273
        if (input_args != "") {
          input_args += ",";
        }
        input_args += out_type;
        input_args += out_name;
274
        input_args_num++;
L
Leo Chen 已提交
275 276 277 278 279

        if (output.dispensable()) {
          const auto out_template =
              output.duplicable() ? OUTPUT_INITIALIZER_TEMPLATE_WITH_NULL_LIST
                                  : OUTPUT_INITIALIZER_TEMPLATE_WITH_NULL;
280 281
          outs_initializer_with_null +=
              paddle::string::Sprintf(out_template, out_name, out_name);
L
Leo Chen 已提交
282 283 284 285 286 287 288 289
        } else {
          const auto out_template = output.duplicable()
                                        ? INPUT_LIST_INITIALIZER_TEMPLATE
                                        : INPUT_INITIALIZER_TEMPLATE;
          outs_initializer +=
              paddle::string::Sprintf(out_template, out_name, out_name);
          outs_initializer += ",";
        }
290 291 292 293 294 295 296 297 298 299 300
      } else {
        // There are few Operators that have duplicable output, like `Out` in
        // split op. We need to specify the number of variables for the
        // duplicable output, as the argument OutNum;
        if (output.duplicable()) {
          if (input_args != "") {
            input_args += ",";
          }
          auto out_num_str = paddle::string::Sprintf(ARG_OUT_NUM, out_name);
          input_args += ARG_OUT_NUM_TYPE;
          input_args += out_num_str;
301
          input_args_num++;
L
Leo Chen 已提交
302
          outs_initializer += paddle::string::Sprintf(
303 304
              OUT_DUPLICABLE_INITIALIZER_TEMPLATE, out_name, out_num_str);
        } else {
L
Leo Chen 已提交
305
          outs_initializer +=
306 307
              paddle::string::Sprintf(OUT_INITIALIZER_TEMPLATE, out_name);
        }
L
Leo Chen 已提交
308
        outs_initializer += ",";
309 310 311 312 313 314 315
      }

      return_type += out_type;
      return_type += ",";
      return_str += paddle::string::Sprintf(return_template, out_name);
      return_str += ",";
      outs_num += 1;
316 317 318
    }
    if (outs_initializer.back() == ',') {
      outs_initializer.pop_back();
319 320
      return_type.pop_back();
      return_str.pop_back();
321 322
    }
    outs_initializer += "}";
323 324 325 326 327 328 329 330 331
    if (outs_num == 0) {
      return_type = "void";
    }
    if (outs_num > 1) {
      return_str = paddle::string::Sprintf(RETURN_TUPLE_TEMPLATE, return_str);
      return_type = paddle::string::Sprintf(RETURN_TUPLE_TYPE, return_type);
    }
    std::string function_args = "";
    if (input_args == "") {
332
      function_args = FUNCTION_ARGS_NO_INPUT;
333 334 335
    } else {
      function_args = paddle::string::Sprintf(FUNCTION_ARGS, input_args);
    }
336

337
    std::string func_name = "imperative_" + op_type;
338
    // generate op funtcion body
339
    auto op_function_str = paddle::string::Sprintf(
340
        OP_FUNCTION_TEMPLATE, return_type, func_name, function_args,
341 342 343
        ins_cast_str, op_type, input_args_num, outs_initializer,
        ins_initializer, ins_initializer_with_null + outs_initializer_with_null,
        op_type, return_str);
344 345

    // generate pybind item
346 347 348 349 350
    auto bind_function_str = paddle::string::Sprintf(
        PYBIND_ITEM_TEMPLATE, module_name, op_type, func_name);

    op_function_list.emplace_back(std::move(op_function_str));
    bind_function_list.emplace_back(std::move(bind_function_str));
351
  }
352
  return std::make_tuple(op_function_list, bind_function_list);
353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370
}

int main(int argc, char* argv[]) {
  if (argc != 2) {
    std::cerr << "argc must be 2" << std::endl;
    return -1;
  }

  std::vector<std::string> headers{"\"paddle/fluid/imperative/tracer.h\""};

  std::ofstream out(argv[1], std::ios::out);

  out << "#pragma once\n\n";

  for (auto& header : headers) {
    out << "#include  " + header + "\n";
  }

371 372
  auto op_funcs = GenerateOpFunctions("m");

373 374 375
  out << "namespace py = pybind11;"
      << "\n";
  out << "namespace paddle {\n"
376 377 378
      << "namespace pybind {\n";
  out << paddle::string::join_strings(std::get<0>(op_funcs), '\n');
  out << "\n\n";
379

380 381
  out << "inline void BindOpFunctions(pybind11::module *module) {\n"
      << "  auto m = module->def_submodule(\"ops\");\n\n";
382

383 384
  out << paddle::string::join_strings(std::get<1>(op_funcs), '\n');
  out << "\n";
385 386 387 388 389 390 391
  out << "}\n\n"
      << "} // namespace pybind\n"
      << "} // namespace paddle\n";

  out.close();
  return 0;
}