linkers_mpi.cpp 1.01 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
#ifdef USE_MPI
6

Guolin Ke's avatar
Guolin Ke committed
7
8
9
10
#include "linkers.h"

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
11
Linkers::Linkers(Config) {
12
  is_init_ = false;
Guolin Ke's avatar
Guolin Ke committed
13
14
15
16
17
18
19
20
21
22
23
24
25
  int argc = 0;
  char**argv = nullptr;
  int flag = 0;
  MPI_SAFE_CALL(MPI_Initialized(&flag));  // test if MPI has been initialized
  if (!flag) {  // if MPI not started, start it
    MPI_SAFE_CALL(MPI_Init_thread(&argc, &argv, MPI_THREAD_SERIALIZED, &flag));
  }
  MPI_SAFE_CALL(MPI_Comm_size(MPI_COMM_WORLD, &num_machines_));
  MPI_SAFE_CALL(MPI_Comm_rank(MPI_COMM_WORLD, &rank_));
  // wait for all client start up
  MPI_SAFE_CALL(MPI_Barrier(MPI_COMM_WORLD));
  bruck_map_ = BruckMap::Construct(rank_, num_machines_);
  recursive_halving_map_ = RecursiveHalvingMap::Construct(rank_, num_machines_);
26
  is_init_ = true;
Guolin Ke's avatar
Guolin Ke committed
27
28
29
}

Linkers::~Linkers() {
30
31
32
  if (is_init_) {
    MPI_SAFE_CALL(MPI_Finalize());
  }
Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
}


}  // namespace LightGBM
37
#endif  // USE_MPI