ipu_compiler.cc 22.7 KB
Newer Older
J
jianghaicheng 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// 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.

A
Allen Guo 已提交
15
#include "paddle/fluid/platform/device/ipu/ipu_compiler.h"
J
jianghaicheng 已提交
16

A
Allen Guo 已提交
17 18 19 20
#include <popart/adam.hpp>
#include <popart/adaptive.hpp>
#include <popart/optimizer.hpp>
#include <popart/sgd.hpp>
J
jianghaicheng 已提交
21
#include "paddle/fluid/framework/ir/graph_helper.h"
A
Allen Guo 已提交
22
#include "paddle/fluid/platform/device/ipu/ipu_utils.h"
J
jianghaicheng 已提交
23 24 25 26 27

namespace paddle {
namespace platform {
namespace ipu {

A
Allen Guo 已提交
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
popart::AdamMode AdamModeFromStr(const std::string& str) {
  if (str == "adam") {
    return popart::AdamMode::Adam;
  } else if (str == "adamax") {
    return popart::AdamMode::AdaMax;
  } else if (str == "lamb") {
    return popart::AdamMode::Lamb;
  } else {
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Uknown AdamMode: %s, AdamMode must be one of these values: adam, "
        "adamax or lamb",
        str));
  }
}

popart::AdaptiveMode AdaptiveModeFromStr(const std::string& str) {
  if (str == "adadelta") {
    return popart::AdaptiveMode::AdaDelta;
  } else if (str == "adagrad") {
    return popart::AdaptiveMode::AdaGrad;
  } else if (str == "rmsprop") {
    return popart::AdaptiveMode::RMSProp;
  } else if (str == "centered_rmsprop") {
    return popart::AdaptiveMode::CenteredRMSProp;
  } else {
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Uknown AdaptiveMode: %s, AdaptiveMode must be one of these values: "
        "adadelta, adagrad, rmsprop or centered_rmsprop",
        str));
  }
}

popart::WeightDecayMode WeightDecayModeFromStr(const std::string& str) {
  if (str == "decay") {
    return popart::WeightDecayMode::Decay;
  } else if (str == "l2_regularization") {
    return popart::WeightDecayMode::L2Regularization;
  } else {
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Uknown WeightDecayMode: %s, WeightDecayMode must be decay or "
        "l2_regularization",
        str));
  }
}

J
jianghaicheng 已提交
73
template <typename T>
A
Allen Guo 已提交
74
T GetAttrAllowNull(std::string attr, OpDesc* op_desc) {
J
jianghaicheng 已提交
75 76 77 78 79 80 81 82
  if (op_desc->HasAttr(attr)) {
    return BOOST_GET_CONST(T, op_desc->GetAttr(attr));
  } else {
    return {};
  }
}

template <typename T>
A
Allen Guo 已提交
83
nonstd::optional<T> GetOptAttrAllowNull(std::string attr, OpDesc* op_desc) {
J
jianghaicheng 已提交
84 85 86 87 88 89 90
  if (op_desc->HasAttr(attr)) {
    return BOOST_GET_CONST(T, op_desc->GetAttr(attr));
  } else {
    return {};
  }
}

A
Allen Guo 已提交
91 92 93 94 95 96 97 98 99 100
template <typename TI, typename TO>
TO GetCastSigAttrAllowNull(std::string attr, OpDesc* op_desc) {
  if (op_desc->HasAttr(attr)) {
    auto x = BOOST_GET_CONST(TI, op_desc->GetAttr(attr));
    return static_cast<TO>(x);
  } else {
    return {};
  }
}

A
Allen Guo 已提交
101 102 103 104 105 106 107 108 109 110 111 112 113
GraphHelper::GraphHelper(const Graph* g) {
  graph = g;
  sorted_ops = framework::ir::TopologySortOperations(*g);
  for (auto* node : g->Nodes()) {
    nodes_id_map[node->id()] = node;
    if (node->IsVar()) {
      vars_name_map[node->Name()] = node;
      sorted_vars_id.push_back(node->id());
    }
  }
  std::sort(sorted_vars_id.begin(), sorted_vars_id.end());
}

A
Allen Guo 已提交
114 115 116 117 118
Compiler::Compiler() { RegisterOpFunc(); }

Compiler::~Compiler() {
  builder_.reset();
  resources_.reset();
J
jianghaicheng 已提交
119 120
}

A
Allen Guo 已提交
121
void Compiler::Prepare(const Graph* graph) {
A
Allen Guo 已提交
122 123
  builder_ = popart::Builder::create();
  resources_ = std::make_unique<CompilerResources>();
A
Allen Guo 已提交
124
  graph_helper_ = std::make_unique<GraphHelper>(graph);
A
Allen Guo 已提交
125
}
J
jianghaicheng 已提交
126 127 128 129

void Compiler::RegisterOpFunc() {
  VLOG(10) << "enter Compiler::RegisterOpFunc";
#define INT_VEC std::vector<std::int64_t>
A
Allen Guo 已提交
130
#define INT32_VEC std::vector<std::int32_t>
J
jianghaicheng 已提交
131 132 133
#define FLOAT_VEC std::vector<float>
#define FLOAT float
#define INT std::int64_t
A
Allen Guo 已提交
134
#define INT32 std::int32_t
J
jianghaicheng 已提交
135 136 137 138 139 140 141
#define BOOL bool
#define STRING std::string
#define STRING_VEC std::vector<std::string*>
#define NONE

#define ARG(Type, Name) , GetAttrAllowNull<Type>(#Name, op_desc)
#define OPT_ARG(Type, Name) , GetOptAttrAllowNull<Type>(#Name, op_desc)
A
Allen Guo 已提交
142
#define SIG_ARG(TI, TO, Name) , GetCastSigAttrAllowNull<TI, TO>(#Name, op_desc)
J
jianghaicheng 已提交
143 144 145 146 147 148 149
#define POPART_CONST_ARG(Name) , const PopartConstant& Name
#define HOST_SIDE_CONST_ARG(Name) , const HostSideConstant& Name
#define POPART_ATTRIB_VEC_ARG(Name)
#define BODY_ARG(Name) NONE

  name_function_ = {
#define OP_DECL(FuncName, OnnxImpl, Args)                     \
A
Allen Guo 已提交
150
  {#FuncName, [&](OpDesc* op_desc) {                          \
J
jianghaicheng 已提交
151 152 153 154 155 156 157 158 159
     auto op_type = op_desc->Type();                          \
     VLOG(10) << "build op:" << op_type << " args " << #Args; \
     auto inputs = GetOpInputs(op_desc);                      \
     auto output_names = GetOpOutputs(op_desc);               \
     auto debug_context = BuildDebugContext(op_desc);         \
     auto aiGraphcoreOpset = builder_->aiGraphcoreOpset1();   \
     auto aiOnnxOpset = builder_->aiOnnxOpset11();            \
     auto output_ids = OnnxImpl(inputs Args, debug_context);  \
     SetIpuIndexStage(output_ids, op_desc);                   \
A
Allen Guo 已提交
160 161
     SetAMPAttributes(output_ids, op_desc);                   \
     SetSerializeAttributes(output_ids, op_desc);             \
J
jianghaicheng 已提交
162 163
     InsertTensors(output_names, output_ids);                 \
   }},  // NOLINT
A
Allen Guo 已提交
164 165
#include "paddle/fluid/platform/device/ipu/supported_ops_autogen.h"
#include "paddle/fluid/platform/device/ipu/supported_ops_custom.h"
J
jianghaicheng 已提交
166 167 168 169 170 171 172
  };

#undef OP_DECL
#undef BODY_ARG
#undef POPART_ATTRIB_VEC_ARG
#undef HOST_SIDE_CONST_ARG
#undef POPART_CONST_ARG
A
Allen Guo 已提交
173
#undef SIG_ARG
J
jianghaicheng 已提交
174 175 176 177 178 179
#undef OPT_ARG
#undef ARG
#undef NONE
#undef STRING_VEC
#undef STRING
#undef BOOL
A
Allen Guo 已提交
180
#undef INT32
J
jianghaicheng 已提交
181 182 183
#undef INT
#undef FLOAT
#undef FLOAT_VEC
A
Allen Guo 已提交
184
#undef INT32_VEC
J
jianghaicheng 已提交
185 186 187
#undef INT_VEC
}

A
Allen Guo 已提交
188
void Compiler::InitInputs(const std::vector<std::string>& feed_list) {
J
jianghaicheng 已提交
189
  for (const auto& feed_name : feed_list) {
A
Allen Guo 已提交
190 191 192 193 194 195 196 197 198 199 200
    auto* node = graph_helper_->vars_name_map[feed_name];
    auto* var_desc = node->Var();
    VLOG(10) << "feed_name= " << var_desc->Name();
    auto data_type = VarType2PopartType(var_desc->GetDataType());
    popart::TensorInfo input_info{data_type, var_desc->GetShape()};
    VLOG(10) << "popart input_info = " << input_info;
    popart::TensorId tensor_id =
        builder_->addInputTensor(input_info, feed_name);
    VLOG(10) << "popart input tensor id = " << tensor_id;
    resources_->inputs.push_back(tensor_id);
    resources_->tensors.emplace(var_desc->Name(), tensor_id);
J
jianghaicheng 已提交
201 202 203 204 205
  }
}

void Compiler::InitOutputs(const std::vector<std::string>& fetch_list) {
  for (const auto& fetch_name : fetch_list) {
A
Allen Guo 已提交
206 207 208 209 210 211
    auto tensor = resources_->tensors.find(fetch_name);
    PADDLE_ENFORCE_NE(
        tensor, resources_->tensors.end(),
        platform::errors::NotFound(
            "Output tensor %s is not found, please check the model.",
            fetch_name));
J
jianghaicheng 已提交
212 213 214
    VLOG(10) << "fetch_name= " << fetch_name;
    VLOG(10) << "popart output tensor id = " << tensor->second;
    builder_->addOutputTensor(tensor->second);
A
Allen Guo 已提交
215 216 217 218
    resources_->outputs.push_back(tensor->second);
  }
}

A
Allen Guo 已提交
219
void Compiler::LowerConstants(const Scope* scope) {
A
Allen Guo 已提交
220 221
  auto& kid_scope = scope->NewScope();
  VLOG(10) << "enter Compiler::LowerConstants";
A
Allen Guo 已提交
222
  for (auto* node : graph_helper_->sorted_ops) {
A
Allen Guo 已提交
223 224 225 226 227 228 229 230 231 232 233 234 235 236
    auto* op_desc = node->Op();
    auto op_type = op_desc->Type();
    if (op_type == "popart_constant") {
      auto shape =
          BOOST_GET_CONST(std::vector<int64_t>, op_desc->GetAttr("dims"));
      auto dtype_ = BOOST_GET_CONST(int, op_desc->GetAttr("dtype"));
      auto dtype = PopartType2VarType(OnnxDtype2PopartType(dtype_));
      auto tensor_name = op_desc->Output("__outputs__")[0];
      auto* var = kid_scope.Var(tensor_name);
      VLOG(10) << "lowering constant: " << tensor_name;
      auto* tensor = var->GetMutable<framework::LoDTensor>();
      ConstantOpAttrVisitor visitor(tensor, dtype);
      auto value = op_desc->GetAttr("value");
      boost::apply_visitor(visitor, value);
237
      auto ddim = phi::make_ddim(shape);
A
Allen Guo 已提交
238 239 240
      tensor->Resize(ddim);

      auto const_data = std::unique_ptr<popart::ConstVoidData>();
A
Allen Guo 已提交
241 242
      popart::TensorInfo tensor_info(PdDataType2PopartType(tensor->dtype()),
                                     shape);
A
Allen Guo 已提交
243 244 245 246 247
      const_data.reset(new popart::ConstVoidData(tensor->data(), tensor_info));
      popart::TensorId result = builder_->aiOnnxOpset11().constant(*const_data);
      SetIpuIndexStage(result, op_desc);
      resources_->tensors.emplace(tensor_name, result);
    }
J
jianghaicheng 已提交
248
  }
A
Allen Guo 已提交
249
  VLOG(10) << "leave Compiler::LowerConstants";
J
jianghaicheng 已提交
250 251
}

A
Allen Guo 已提交
252
void Compiler::LowerWeights(const Scope* scope) {
A
Allen Guo 已提交
253
  VLOG(10) << "enter Compiler::LowerWeights";
J
jianghaicheng 已提交
254
  // at this step, the graph doesn't contains optimizer related states
A
Allen Guo 已提交
255 256
  for (auto id : graph_helper_->sorted_vars_id) {
    auto* node = graph_helper_->nodes_id_map[id];
J
jianghaicheng 已提交
257 258 259
    if (node->IsVar() && !node->IsCtrlVar() && node->Var()) {
      if (node->Var()->Persistable() && node->inputs.empty()) {
        auto var_name = node->Var()->Name();
A
Allen Guo 已提交
260
        if (resources_->tensors.count(var_name) != 0) {
A
Allen Guo 已提交
261
          VLOG(10) << "found existed one, skip lowering Weight: " << var_name;
J
jianghaicheng 已提交
262 263
          continue;
        }
A
Allen Guo 已提交
264
        VLOG(10) << "lowering weight: " << var_name;
J
jianghaicheng 已提交
265

A
Allen Guo 已提交
266
        auto var = scope->FindVar(var_name);
J
jianghaicheng 已提交
267 268
        if (var) {
          auto tensor = var->Get<framework::LoDTensor>();
A
Allen Guo 已提交
269
          auto dtype = PdDataType2PopartType(tensor.dtype());
J
jianghaicheng 已提交
270 271 272 273 274
          auto shape = std::vector<int64_t>();
          for (size_t i = 0; i < tensor.dims().size(); ++i) {
            shape.push_back(tensor.dims().at(i));
          }
          popart::TensorInfo tensor_info(dtype, shape);
275
          popart::ConstVoidData const_data{tensor.data(), tensor_info};
J
jianghaicheng 已提交
276 277
          popart::TensorId result =
              builder_->addInitializedInputTensor(const_data, var_name);
A
Allen Guo 已提交
278 279
          resources_->tensors.emplace(var_name, result);
          resources_->weights.push_back(result);
J
jianghaicheng 已提交
280 281 282 283
        }
      }
    }
  }
A
Allen Guo 已提交
284 285 286
  VLOG(10) << "leave Compiler::LowerWeights";
}

A
Allen Guo 已提交
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 335 336 337 338 339 340 341
void Compiler::LowerBody() {
  VLOG(10) << "enter Compiler::LowerBody";
  for (auto* node : graph_helper_->sorted_ops) {
    auto* op_desc = node->Op();
    auto op_type = op_desc->Type();
    VLOG(10) << "lowering op: " << op_type;

    if (op_type == "popart_constant") {
      // pass
    } else if (op_type == "popart_optimizer") {
      // pass
    } else if (op_type == "popart_checkpointoutput") {
      auto inputs = GetOpInputs(op_desc);
      auto outputs = GetOpOutputs(op_desc);
      auto output_ids = builder_->checkpointOutput(inputs);
      InsertTensors(outputs, output_ids);
    } else if (op_type == "popart_custom_op") {
      auto inputs = GetOpInputs(op_desc);
      auto outputs = GetOpOutputs(op_desc);
      auto debug_context = BuildDebugContext(op_desc);
      auto attributes = std::map<std::string, popart::any>{};
      for (auto& attr : op_desc->GetAttrMap()) {
        CustomOpAttrVisitor visitor(&attributes, attr.first);
        boost::apply_visitor(visitor, attr.second);
      }
      auto __op_type =
          BOOST_GET_CONST(std::string, op_desc->GetAttr("__op_type"));
      VLOG(10) << "Build graph from custom op: " << __op_type;
      auto it = custom_ops_.find(__op_type);
      auto output_ids =
          builder_->customOp(it->second.popart_op, it->second.popart_op.version,
                             inputs, outputs.size(), attributes, debug_context);
      SetIpuIndexStage(output_ids, op_desc);
      InsertTensors(outputs, output_ids);
    } else if (op_type == "popart_printtensor") {
      auto inputs = GetOpInputs(op_desc);
      auto outputs = GetOpOutputs(op_desc);
      auto debug_context = BuildDebugContext(op_desc);
      auto print_gradient =
          BOOST_GET_CONST(int64_t, op_desc->GetAttr("print_gradient"));
      auto title = BOOST_GET_CONST(std::string, op_desc->GetAttr("title"));
      auto output_ids = builder_->aiGraphcoreOpset1().printtensor(
          inputs, print_gradient, debug_context, title);
      SetIpuIndexStage(output_ids, op_desc);
      InsertTensors(outputs, output_ids);
    } else {
      auto itr = name_function_.find(op_type);
      if (itr != name_function_.end()) {
        itr->second(node->Op());
      } else {
        PADDLE_THROW(platform::errors::NotFound(
            "%s is not registered, please check for unsupported operators for "
            "running on IPU",
            op_type));
      }
A
Allen Guo 已提交
342
    }
A
Allen Guo 已提交
343 344 345
  }
  VLOG(10) << "leave Compiler::LowerBody";
}
A
Allen Guo 已提交
346

A
Allen Guo 已提交
347 348
void Compiler::LowerOptimizer(const Scope* scope) {
  for (auto* node : graph_helper_->sorted_ops) {
A
Allen Guo 已提交
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 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440
    auto* op_desc = node->Op();
    auto op_type = op_desc->Type();
    if (op_type == "popart_optimizer") {
      auto raw_type =
          BOOST_GET_CONST(std::string, op_desc->GetAttr("raw_type"));
      resources_->optimizer_type = raw_type;
      auto loss_var =
          BOOST_GET_CONST(std::string, op_desc->GetAttr("loss_var"));
      resources_->loss_var = resources_->tensors[loss_var];
      resources_->with_lr_sched =
          BOOST_GET_CONST(bool, op_desc->GetAttr("with_lr_sched"));
      if (op_desc->HasAttr("lr_var")) {
        auto lr_var = BOOST_GET_CONST(std::string, op_desc->GetAttr("lr_var"));
        resources_->lr_var = lr_var;
        resources_->lr = GetSingleVarFromScope<float>(scope, lr_var);
      } else {
        // adadelta has no lr
        resources_->lr = 0.01f;
        resources_->with_lr_sched = false;
      }
      VLOG(10) << "Set initial lr: " << resources_->lr;
      auto loss_scaling = ipu_strategy_->loss_scaling;
      auto type = BOOST_GET_CONST(std::string, op_desc->GetAttr("type"));
      if (type == "sgd") {
        auto weight_decay =
            BOOST_GET_CONST(float, op_desc->GetAttr("weight_decay"));
        auto momentum = BOOST_GET_CONST(float, op_desc->GetAttr("momentum"));
        resources_->optimizer_fn = [=](float lr) {
          return std::make_unique<popart::SGD>(
              popart::OptimizerValue(lr, false),
              popart::OptimizerValue(weight_decay, true),
              popart::OptimizerValue(momentum, true),
              popart::SGD::getUnsetDampening(),
              popart::SGD::getUnsetVelocityScaling(),
              popart::OptimizerValue(loss_scaling, true));
        };
      } else if (type == "adam") {
        auto weight_decay =
            BOOST_GET_CONST(float, op_desc->GetAttr("weight_decay"));
        auto beta1 = BOOST_GET_CONST(float, op_desc->GetAttr("beta1"));
        auto beta2 = BOOST_GET_CONST(float, op_desc->GetAttr("beta2"));
        auto eps = BOOST_GET_CONST(float, op_desc->GetAttr("eps"));
        auto mwn = ipu_strategy_->max_weight_norm;
        VLOG(10) << "set max_weight_norm: " << mwn;
        auto adam_mode_ =
            BOOST_GET_CONST(std::string, op_desc->GetAttr("adam_mode"));
        auto adam_mode = AdamModeFromStr(adam_mode_);
        auto weight_decay_mode_ =
            BOOST_GET_CONST(std::string, op_desc->GetAttr("weight_decay_mode"));
        auto weight_decay_mode = WeightDecayModeFromStr(weight_decay_mode_);
        resources_->optimizer_fn = [=](float lr) {
          return std::make_unique<popart::Adam>(
              popart::OptimizerValue(lr, false),
              popart::OptimizerValue(weight_decay, true),
              popart::OptimizerValue(beta1, true),
              popart::OptimizerValue(beta2, true),
              popart::OptimizerValue(eps, true),
              popart::OptimizerValue(loss_scaling, true),
              popart::OptimizerValue(mwn, true), adam_mode, weight_decay_mode,
              popart::DataType::UNDEFINED, popart::DataType::FLOAT,
              popart::DataType::FLOAT);
        };
      } else if (type == "adaptive") {
        auto alpha = BOOST_GET_CONST(float, op_desc->GetAttr("alpha"));
        auto momentum = BOOST_GET_CONST(float, op_desc->GetAttr("momentum"));
        auto eps = BOOST_GET_CONST(float, op_desc->GetAttr("eps"));
        auto weight_decay =
            BOOST_GET_CONST(float, op_desc->GetAttr("weight_decay"));
        auto adaptive_mode_ =
            BOOST_GET_CONST(std::string, op_desc->GetAttr("adaptive_mode"));
        auto adaptive_mode = AdaptiveModeFromStr(adaptive_mode_);
        auto weight_decay_mode_ =
            BOOST_GET_CONST(std::string, op_desc->GetAttr("weight_decay_mode"));
        auto weight_decay_mode = WeightDecayModeFromStr(weight_decay_mode_);
        resources_->optimizer_fn = [=](float lr) {
          return std::make_unique<popart::Adaptive>(
              popart::OptimizerValue(lr, false),
              popart::OptimizerValue(weight_decay, true),
              popart::OptimizerValue(alpha, true),
              popart::OptimizerValue(momentum, true),
              popart::OptimizerValue(eps, true),
              popart::OptimizerValue(loss_scaling, true), adaptive_mode,
              weight_decay_mode, popart::DataType::UNDEFINED,
              popart::DataType::FLOAT, popart::DataType::FLOAT,
              popart::DataType::FLOAT);
        };
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "optimizer %s is not implemented", type));
      }
    }
  }
J
jianghaicheng 已提交
441 442 443 444 445 446 447 448
}

void Compiler::InsertTensors(const std::vector<std::string>& output_names,
                             const std::vector<std::string>& tensor_ids) {
  PADDLE_ENFORCE_EQ(output_names.size(), tensor_ids.size(),
                    platform::errors::Fatal("InsertTensors size mismatch"));
  for (int i = 0; i < tensor_ids.size(); i++) {
    std::string tensor_id = tensor_ids[i];
A
Allen Guo 已提交
449
    resources_->tensors.emplace(output_names[i], tensor_ids[i]);
J
jianghaicheng 已提交
450 451 452 453 454 455 456
  }
}

void Compiler::InsertTensors(const std::vector<std::string>& output_names,
                             const std::string& tensor_id) {
  PADDLE_ENFORCE_EQ(output_names.size(), 1,
                    platform::errors::Fatal("InsertTensors size mismatch"));
A
Allen Guo 已提交
457
  resources_->tensors.emplace(output_names[0], tensor_id);
J
jianghaicheng 已提交
458 459 460
}

void Compiler::SetIpuIndexStage(const std::vector<std::string>& tensor_ids,
A
Allen Guo 已提交
461
                                const OpDesc* op_desc) {
J
jianghaicheng 已提交
462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481
  VLOG(10) << "enter Compiler::SetIpuIndexStage";
  auto tensor_ids_set =
      std::set<std::string>(tensor_ids.begin(), tensor_ids.end());

  if (op_desc->HasAttr(sIpuIndexAttr)) {
    auto ipu_index = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuIndexAttr));
    builder_->virtualGraph(tensor_ids_set, ipu_index);
    VLOG(10) << "set " << sIpuIndexAttr << " = " << ipu_index
             << " for op: " << op_desc->Type();
    if (op_desc->HasAttr(sIpuStageAttr)) {
      auto ipu_stage = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuStageAttr));
      builder_->pipelineStage(tensor_ids_set, ipu_stage);
      VLOG(10) << "set " << sIpuStageAttr << "= " << ipu_stage
               << " for op: " << op_desc->Type();
    }
  }
  VLOG(10) << "leave Compiler::SetIpuIndexStage";
}

void Compiler::SetIpuIndexStage(const std::string& tensor_id,
A
Allen Guo 已提交
482
                                const OpDesc* op_desc) {
J
jianghaicheng 已提交
483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
  VLOG(10) << "enter Compiler::SetIpuIndexStage";

  if (op_desc->HasAttr(sIpuIndexAttr)) {
    auto ipu_index = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuIndexAttr));
    builder_->virtualGraph(tensor_id, ipu_index);
    VLOG(10) << "set " << sIpuIndexAttr << " = " << ipu_index
             << " for op: " << op_desc->Type();
    if (op_desc->HasAttr(sIpuStageAttr)) {
      auto ipu_stage = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuStageAttr));
      builder_->pipelineStage(tensor_id, ipu_stage);
      VLOG(10) << "set " << sIpuStageAttr << "= " << ipu_stage
               << " for op: " << op_desc->Type();
    }
  }
  VLOG(10) << "leave Compiler::SetIpuIndexStage";
}

A
Allen Guo 已提交
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 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546
void Compiler::SetAMPAttributes(const std::vector<std::string>& tensor_ids,
                                const OpDesc* op_desc) {
  if (op_desc->Type() == "popart_matmul") {
    for (const auto& tensor_id : tensor_ids) {
      SetAMPAttributes(tensor_id, op_desc);
    }
  }
}

void Compiler::SetAMPAttributes(const std::string& tensor_id,
                                const OpDesc* op_desc) {
  VLOG(10) << "enter Compiler::SetAMPAttributes";
  if (op_desc->Type() == "popart_matmul") {
    auto amp = ipu_strategy_->available_memory_proportion;
    if (amp > 0.0f && amp <= 1.0) {
      builder_->setAvailableMemoryProportion(tensor_id, amp);
    }
  }
  VLOG(10) << "leave Compiler::SetAMPAttributes";
}

void Compiler::SetSerializeAttributes(
    const std::vector<std::string>& tensor_ids, const OpDesc* op_desc) {
  VLOG(10) << "enter Compiler::SetSerializeAttributes";
  auto tensor_ids_set =
      std::set<std::string>(tensor_ids.begin(), tensor_ids.end());

  if (op_desc->Type() == "popart_matmul") {
    if (op_desc->HasAttr(sMatmulSerializeFactor)) {
      auto factor =
          BOOST_GET_CONST(int, op_desc->GetAttr(sMatmulSerializeFactor));
      std::string mode = "output_channels";
      if (op_desc->HasAttr(sMatmulSerializeMode)) {
        mode = BOOST_GET_CONST(std::string,
                               op_desc->GetAttr(sMatmulSerializeMode));
      }
      builder_->setSerializeMatMul(tensor_ids_set, mode, (int64_t)factor, true);
    }
  }
  VLOG(10) << "leave Compiler::SetSerializeAttributes";
}

void Compiler::SetSerializeAttributes(const std::string& tensor_id,
                                      const OpDesc* op_desc) {
  std::vector<std::string> tensor_ids = {tensor_id};
  SetSerializeAttributes(tensor_ids, op_desc);
}
J
jianghaicheng 已提交
547

A
Allen Guo 已提交
548 549 550 551 552 553 554 555
void Compiler::SetCustomOps(
    const std::vector<IpuCustomOpIdentifier>& custom_ops) {
  for (auto x : custom_ops) {
    custom_ops_.emplace(x.paddle_op, x);
  }
}

std::string Compiler::GetFP16ModelProto() {
J
jianghaicheng 已提交
556 557
  popart::GraphTransformer graph_transformer(builder_->getModelProto());
  graph_transformer.convertFloatsToHalfs();
A
Allen Guo 已提交
558
  return graph_transformer.getModelProto();
J
jianghaicheng 已提交
559 560 561
}

std::string Compiler::GetModelProto() {
A
Allen Guo 已提交
562 563 564 565
  if (ipu_strategy_->enable_fp16) {
    return GetFP16ModelProto();
  } else {
    return builder_->getModelProto();
J
jianghaicheng 已提交
566 567 568 569 570 571 572 573 574 575 576 577 578 579
  }
}

void Compiler::SaveModelProto(const std::string& path) {
  builder_->saveModelProto(path);
}

void Compiler::SaveModelProtoNoCheck(const std::string& path) {
  auto proto = GetModelProto();
  std::ofstream onnxfile(path, std::ios_base::binary);
  onnxfile.write(proto.data(), proto.size());
  onnxfile.close();
}

A
Allen Guo 已提交
580
std::vector<std::string> Compiler::GetOpInputs(const OpDesc* op) {
J
jianghaicheng 已提交
581 582 583
  auto ins = op->Input("__inputs__");
  std::vector<std::string> inputs;
  for (const auto& in : ins) {
A
Allen Guo 已提交
584 585
    if (resources_->tensors.find(in) != resources_->tensors.end()) {
      inputs.push_back(resources_->tensors[in]);
J
jianghaicheng 已提交
586 587 588 589 590 591 592
    } else {
      inputs.push_back(in);
    }
  }
  return inputs;
}

A
Allen Guo 已提交
593
const std::vector<std::string>& Compiler::GetOpOutputs(const OpDesc* op) {
J
jianghaicheng 已提交
594 595 596
  return op->Output("__outputs__");
}

A
Allen Guo 已提交
597
popart::DebugContext Compiler::BuildDebugContext(const OpDesc* op) {
J
jianghaicheng 已提交
598 599 600 601 602 603 604 605 606 607
  auto op_identify_id =
      BOOST_GET_CONST(std::string, op->GetAttr(sOpIdentifyIdAttr));
  VLOG(10) << "op_identify_id of op: " << op->Type() << " is "
           << op_identify_id;
  return popart::DebugContext(op_identify_id);
}

}  // namespace ipu
}  // namespace platform
}  // namespace paddle