Merge pull request #388 from ValeevGroup/evaleev/fix/ta-initialize

[python] can load with initialized MPI

python/src/TiledArray/python/module.h

@@ -17,82 +17,98 @@
  *
  */
 
+#include "array.h"
+#include "einsum.h"
+#include "expression.h"
 #include "python.h"
 #include "range.h"
 #include "trange.h"
-#include "array.h"
-#include "expression.h"
-#include "einsum.h"
 
-#include <tiledarray.h>
 #include <dlfcn.h>
+#include <tiledarray.h>
 
 namespace TiledArray {
 namespace python {
 
-  static World& initialize() {
+static bool initialized_madness = false;
 
-#if defined( __linux__) && defined(OPEN_MPI)
-    dlopen("libmpi.so", RTLD_NOW | RTLD_NOLOAD | RTLD_GLOBAL); // ompi hack
+static World &initialize() {
+#if defined(__linux__) && defined(OPEN_MPI)
+  dlopen("libmpi.so", RTLD_NOW | RTLD_NOLOAD | RTLD_GLOBAL);  // ompi hack
 #endif
 
-    // this loads MPI before TA tries to do it
-    int initialized = 0;
-    MPI_Initialized(&initialized);
-
-    if (!initialized) {
-      // Initialize runtime
-      int argc = 0;
-      char *_argv[0];
-      char **argv = _argv;
-      //MPI_Init(&argc, &argv);
-      TiledArray::World& world = TiledArray::initialize(argc, argv);
-      if (world.rank() == 0) {
-        std::cout << "initialized TA in a world with " << world.size() << " ranks" << std::endl;
-      }
-    }
-
-    return get_default_world();
-
+  // this loads MPI before TA tries to do it
+  int initialized = 0;
+  MPI_Initialized(&initialized);
+  MPI_Comm ta_comm;
+
+  if (!initialized) {
+    ta_comm = MPI_COMM_WORLD;
+  } else {
+    int thread_level;
+    MPI_Query_thread(&thread_level);
+    if (thread_level != MPI_THREAD_MULTIPLE)
+      TA_EXCEPTION(
+          "If initializing MPI before loading PyTA, must use "
+          "MPI_Init_thread(argc,argv,MPI_THREAD_MULTIPLE)");
   }
 
-  void finalize() {
-    if (!TiledArray::finalized()) {
-      TiledArray::finalize();
-    }
+  // Initialize runtime
+  int argc = 0;
+  char *_argv[0];
+  char **argv = _argv;
+  if (!madness::initialized()) {
+    madness::initialize(argc, argv, ta_comm);
+    initialized_madness = true;
   }
+  TiledArray::World &world = TiledArray::initialize(argc, argv, ta_comm);
+  TiledArray::set_default_world(world);
+  if (world.rank() == 0) {
+    std::cout << "initialized TA in a world with " << world.size() << " ranks"
+              << std::endl;
+  }
+
+  return get_default_world();
+}
 
-  py::object default_world() {
-    return py::cast(TiledArray::get_default_world(), py::return_value_policy::reference);
+void finalize() {
+  if (!TiledArray::finalized()) {
+    TiledArray::finalize();
   }
+  if (initialized_madness) {
+    madness::finalize();
+  }
+}
 
-  World& initialize(py::module m) {
+py::object default_world() {
+  return py::cast(TiledArray::get_default_world(),
+                  py::return_value_policy::reference);
+}
 
-    auto &world = initialize();
+World &initialize(py::module m) {
+  auto &world = initialize();
 
-    py::class_< TiledArray::World >(m, "World")
+  py::class_<TiledArray::World>(m, "World")
       // .def_property_readonly_static(
       //   "COMM_WORLD",
       //   [](){ return madness::World::find_instance(SafeMPI::COMM_WORLD); }
       // )
       .def_property_readonly("rank", &TiledArray::World::rank)
       .def_property_readonly("size", &TiledArray::World::size)
-      .def("fence", [](TiledArray::World &world) { world.gop.fence(); })
-      ;
+      .def("fence", [](TiledArray::World &world) { world.gop.fence(); });
 
-    m.def("get_default_world", &default_world);
+  m.def("get_default_world", &default_world);
 
-    m.def("finalize", &TiledArray::python::finalize);
+  m.def("finalize", &TiledArray::python::finalize);
 
-    range::__init__(m);
-    trange::__init__(m);
-    expression::__init__(m);
-    array::__init__(m);
-    einsum::__init__(m);
-
-    return world;
-
-  }
+  range::__init__(m);
+  trange::__init__(m);
+  expression::__init__(m);
+  array::__init__(m);
+  einsum::__init__(m);
 
+  return world;
 }
-}
+
+}  // namespace python
+}  // namespace TiledArray

src/TiledArray/external/madness.h

@@ -54,6 +54,12 @@ struct default_world {
       TA_ASSERT(madness::initialized() &&
                 "TiledArray::detail::default_world::get() called "
                 "before madness::initialize() OR after madness::finalize()");
+#if MADNESS_DISABLE_WORLD_GET_DEFAULT
+      TA_EXCEPTION(
+          "TiledArray::set_default_world() must be called before calling "
+          "TiledArray::get_default_world() if MADWorld configured with "
+          "DISABLE_WORLD_GET_DEFAULT=ON");
+#endif
       world() = &madness::World::get_default();
     }
     return *world();

src/TiledArray/util/backtrace.cpp

@@ -28,6 +28,7 @@
 #include <TiledArray/util/backtrace.h>
 
 #include <cstring>
+#include <iostream>
 #include <iterator>
 #include <sstream>
 
@@ -187,3 +188,8 @@ std::string Backtrace::__demangle(const std::string &symbol) {
 
 }  // namespace detail
 }  // namespace TiledArray
+
+extern "C" void tiledarray_dump_backtrace_to_std_cout() {
+  TiledArray::detail::Backtrace bt("tiledarray_dump_backtrace: ");
+  std::cout << bt.str() << std::endl;
+}

src/TiledArray/util/backtrace.h

@@ -25,8 +25,8 @@
 // The U.S. Government is granted a limited license as per AL 91-7.
 //
 
-#ifndef MPQC4_SRC_MPQC_UTIL_CORE_BACKTRACE_H_
-#define MPQC4_SRC_MPQC_UTIL_CORE_BACKTRACE_H_
+#ifndef TILEDARRAY_UTIL_BACKTRACE_H_
+#define TILEDARRAY_UTIL_BACKTRACE_H_
 
 #include <string>
 #include <vector>
@@ -80,4 +80,8 @@ class Backtrace {
 }  // namespace detail
 }  // namespace TiledArray
 
-#endif  // MPQC4_SRC_MPQC_UTIL_CORE_BACKTRACE_H_
+/// dumps backtrace to std::cout
+/// @note this function has C linkage
+extern "C" void tiledarray_dump_backtrace_to_std_cout();
+
+#endif  // TILEDARRAY_UTIL_BACKTRACE_H_