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 101 102 103 104 105
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 {};
  }
}

Compiler::Compiler() { RegisterOpFunc(); }

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

A
Allen Guo 已提交
108 109 110 111
void Compiler::Prepare() {
  builder_ = popart::Builder::create();
  resources_ = std::make_unique<CompilerResources>();
}
J
jianghaicheng 已提交
112 113 114 115

void Compiler::RegisterOpFunc() {
  VLOG(10) << "enter Compiler::RegisterOpFunc";
#define INT_VEC std::vector<std::int64_t>
A
Allen Guo 已提交
116
#define INT32_VEC std::vector<std::int32_t>
J
jianghaicheng 已提交
117 118 119
#define FLOAT_VEC std::vector<float>
#define FLOAT float
#define INT std::int64_t
A
Allen Guo 已提交
120
#define INT32 std::int32_t
J
jianghaicheng 已提交
121 122 123 124 125 126 127
#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 已提交
128
#define SIG_ARG(TI, TO, Name) , GetCastSigAttrAllowNull<TI, TO>(#Name, op_desc)
J
jianghaicheng 已提交
129 130 131 132 133 134 135
#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 已提交
136
  {#FuncName, [&](OpDesc* op_desc) {                          \
J
jianghaicheng 已提交
137 138 139 140 141 142 143 144 145
     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 已提交
146 147
     SetAMPAttributes(output_ids, op_desc);                   \
     SetSerializeAttributes(output_ids, op_desc);             \
J
jianghaicheng 已提交
148 149
     InsertTensors(output_names, output_ids);                 \
   }},  // NOLINT
A
Allen Guo 已提交
150 151
#include "paddle/fluid/platform/device/ipu/supported_ops_autogen.h"
#include "paddle/fluid/platform/device/ipu/supported_ops_custom.h"
J
jianghaicheng 已提交
152 153 154 155 156 157 158
  };

#undef OP_DECL
#undef BODY_ARG
#undef POPART_ATTRIB_VEC_ARG
#undef HOST_SIDE_CONST_ARG
#undef POPART_CONST_ARG
A
Allen Guo 已提交
159
#undef SIG_ARG
J
jianghaicheng 已提交
160 161 162 163 164 165
#undef OPT_ARG
#undef ARG
#undef NONE
#undef STRING_VEC
#undef STRING
#undef BOOL
A
Allen Guo 已提交
166
#undef INT32
J
jianghaicheng 已提交
167 168 169
#undef INT
#undef FLOAT
#undef FLOAT_VEC
A
Allen Guo 已提交
170
#undef INT32_VEC
J
jianghaicheng 已提交
171 172 173
#undef INT_VEC
}

A
Allen Guo 已提交
174
void Compiler::LowerBody(const Graph* graph) {
J
jianghaicheng 已提交
175 176 177 178 179
  VLOG(10) << "enter Compiler::LowerBody";
  auto nodes = framework::ir::TopologySortOperations(*graph);
  for (auto* node : nodes) {
    auto* op_desc = node->Op();
    auto op_type = op_desc->Type();
A
Allen Guo 已提交
180
    VLOG(10) << "lowering op: " << op_type;
J
jianghaicheng 已提交
181 182

    if (op_type == "popart_constant") {
A
Allen Guo 已提交
183 184 185 186
      // pass
    } else if (op_type == "popart_optimizer") {
      // pass
    } else if (op_type == "popart_checkpointoutput") {
J
jianghaicheng 已提交
187 188
      auto inputs = GetOpInputs(op_desc);
      auto outputs = GetOpOutputs(op_desc);
A
Allen Guo 已提交
189 190 191
      auto output_ids = builder_->checkpointOutput(inputs);
      InsertTensors(outputs, output_ids);
    } else if (op_type == "popart_custom_op") {
J
jianghaicheng 已提交
192 193
      auto inputs = GetOpInputs(op_desc);
      auto outputs = GetOpOutputs(op_desc);
A
Allen Guo 已提交
194 195 196 197 198
      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);
J
jianghaicheng 已提交
199
      }
A
Allen Guo 已提交
200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219
      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);
J
jianghaicheng 已提交
220 221 222 223 224 225
    } else {
      auto itr = name_function_.find(op_type);
      if (itr != name_function_.end()) {
        itr->second(node->Op());
      } else {
        PADDLE_THROW(platform::errors::NotFound(
A
Allen Guo 已提交
226 227 228
            "%s is not registered, please check for unsupported operators for "
            "running on IPU",
            op_type));
J
jianghaicheng 已提交
229 230 231 232 233 234
      }
    }
  }
  VLOG(10) << "leave Compiler::LowerBody";
}

A
Allen Guo 已提交
235
void Compiler::InitInputs(Graph* graph,
J
jianghaicheng 已提交
236 237 238
                          const std::vector<std::string>& feed_list) {
  for (const auto& feed_name : feed_list) {
    feed_list_.push_back(feed_name);
A
Allen Guo 已提交
239
    for (const Node* n : graph->Nodes()) {
J
jianghaicheng 已提交
240 241 242 243 244 245 246 247 248 249
      if (n->IsVar()) {
        auto* var_desc = n->Var();
        if (feed_name == var_desc->Name()) {
          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;
A
Allen Guo 已提交
250 251
          resources_->inputs.push_back(tensor_id);
          resources_->tensors.emplace(var_desc->Name(), tensor_id);
J
jianghaicheng 已提交
252 253 254 255 256 257 258 259 260
        }
      }
    }
  }
}

void Compiler::InitOutputs(const std::vector<std::string>& fetch_list) {
  for (const auto& fetch_name : fetch_list) {
    fetch_list_.push_back(fetch_name);
A
Allen Guo 已提交
261 262 263 264 265 266
    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 已提交
267 268 269
    VLOG(10) << "fetch_name= " << fetch_name;
    VLOG(10) << "popart output tensor id = " << tensor->second;
    builder_->addOutputTensor(tensor->second);
A
Allen Guo 已提交
270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295
    resources_->outputs.push_back(tensor->second);
  }
}

void Compiler::LowerConstants(const Graph* graph, const Scope* scope) {
  auto& kid_scope = scope->NewScope();
  VLOG(10) << "enter Compiler::LowerConstants";
  for (auto* node : graph->Nodes()) {
    if (!node->IsOp()) {
      continue;
    }

    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);
296
      auto ddim = pten::make_ddim(shape);
A
Allen Guo 已提交
297 298 299
      tensor->Resize(ddim);

      auto const_data = std::unique_ptr<popart::ConstVoidData>();
A
Allen Guo 已提交
300 301
      popart::TensorInfo tensor_info(PdDataType2PopartType(tensor->dtype()),
                                     shape);
A
Allen Guo 已提交
302 303 304 305 306
      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 已提交
307
  }
A
Allen Guo 已提交
308
  VLOG(10) << "leave Compiler::LowerConstants";
J
jianghaicheng 已提交
309 310
}

A
Allen Guo 已提交
311 312 313
void Compiler::LowerWeights(const Graph* graph, const Scope* scope) {
  VLOG(10) << "enter Compiler::LowerWeights";
  PADDLE_ENFORCE_NOT_NULL(scope,
J
jianghaicheng 已提交
314 315 316 317 318 319 320
                          platform::errors::PreconditionNotMet(
                              "You should call set_scope before LowerWeights"));
  // at this step, the graph doesn't contains optimizer related states
  for (const auto* node : graph->Nodes()) {
    if (node->IsVar() && !node->IsCtrlVar() && node->Var()) {
      if (node->Var()->Persistable() && node->inputs.empty()) {
        auto var_name = node->Var()->Name();
A
Allen Guo 已提交
321
        if (resources_->tensors.count(var_name) != 0) {
J
jianghaicheng 已提交
322 323
          continue;
        }
A
Allen Guo 已提交
324
        VLOG(10) << "lowering weight: " << var_name;
J
jianghaicheng 已提交
325

A
Allen Guo 已提交
326
        auto var = scope->FindVar(var_name);
J
jianghaicheng 已提交
327 328
        if (var) {
          auto tensor = var->Get<framework::LoDTensor>();
A
Allen Guo 已提交
329
          auto dtype = PdDataType2PopartType(tensor.dtype());
J
jianghaicheng 已提交
330 331 332 333 334
          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);
335
          popart::ConstVoidData const_data{tensor.data(), tensor_info};
J
jianghaicheng 已提交
336 337
          popart::TensorId result =
              builder_->addInitializedInputTensor(const_data, var_name);
A
Allen Guo 已提交
338 339
          resources_->tensors.emplace(var_name, result);
          resources_->weights.push_back(result);
J
jianghaicheng 已提交
340 341 342 343
        }
      }
    }
  }
A
Allen Guo 已提交
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 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 441 442 443 444
  VLOG(10) << "leave Compiler::LowerWeights";
}

void Compiler::LowerOptimier(const Graph* graph, const Scope* scope) {
  for (auto* node : graph->Nodes()) {
    if (!node->IsOp()) {
      continue;
    }

    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 已提交
445 446 447 448 449 450 451 452
}

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 已提交
453
    resources_->tensors.emplace(output_names[i], tensor_ids[i]);
J
jianghaicheng 已提交
454 455 456 457 458 459 460
  }
}

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 已提交
461
  resources_->tensors.emplace(output_names[0], tensor_id);
J
jianghaicheng 已提交
462 463 464
}

void Compiler::SetIpuIndexStage(const std::vector<std::string>& tensor_ids,
A
Allen Guo 已提交
465
                                const OpDesc* op_desc) {
J
jianghaicheng 已提交
466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485
  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 已提交
486
                                const OpDesc* op_desc) {
J
jianghaicheng 已提交
487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503
  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 已提交
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 547 548 549 550
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 已提交
551

A
Allen Guo 已提交
552 553 554 555 556 557 558 559
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 已提交
560 561
  popart::GraphTransformer graph_transformer(builder_->getModelProto());
  graph_transformer.convertFloatsToHalfs();
A
Allen Guo 已提交
562
  return graph_transformer.getModelProto();
J
jianghaicheng 已提交
563 564 565
}

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

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

A
Allen Guo 已提交
597
const std::vector<std::string>& Compiler::GetOpOutputs(const OpDesc* op) {
J
jianghaicheng 已提交
598 599 600
  return op->Output("__outputs__");
}

A
Allen Guo 已提交
601
popart::DebugContext Compiler::BuildDebugContext(const OpDesc* op) {
J
jianghaicheng 已提交
602 603 604 605 606 607 608 609 610 611
  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