Fix tally reduction

include/openmc/capi.h

@@ -3,6 +3,7 @@
 
 #include <stdint.h>
 #include <stdbool.h>
+#include <stddef.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -106,7 +107,7 @@ extern "C" {
   int openmc_tally_get_scores(int32_t index, int** scores, int* n);
   int openmc_tally_get_type(int32_t index, int32_t* type);
   int openmc_tally_reset(int32_t index);
-  int openmc_tally_results(int32_t index, double** ptr, int shape_[3]);
+  int openmc_tally_results(int32_t index, double** ptr, size_t shape_[3]);
   int openmc_tally_set_active(int32_t index, bool active);
   int openmc_tally_set_estimator(int32_t index, const char* estimator);
   int openmc_tally_set_filters(int32_t index, int n, const int32_t* indices);

include/openmc/hdf5_interface.h

@@ -8,6 +8,7 @@
 #include <cstring> // for strlen
 #include <string>
 #include <sstream>
+#include <type_traits>
 #include <vector>
 
 #include "hdf5.h"
@@ -333,7 +334,9 @@ write_attribute(hid_t obj_id, const char* name, const std::vector<T>& buffer)
 // Templates/overloads for write_dataset
 //==============================================================================
 
-template<typename T> inline void
+// Template for scalars (ensured by SFINAE)
+template<typename T> inline
+std::enable_if_t<std::is_scalar<std::decay_t<T>>::value>
 write_dataset(hid_t obj_id, const char* name, T buffer)
 {
   write_dataset(obj_id, 0, nullptr, name, H5TypeMap<T>::type_id, &buffer, false);
@@ -359,18 +362,11 @@ write_dataset(hid_t obj_id, const char* name, const std::vector<T>& buffer)
   write_dataset(obj_id, 1, dims, name, H5TypeMap<T>::type_id, buffer.data(), false);
 }
 
-template<typename T> inline void
-write_dataset(hid_t obj_id, const char* name, const xt::xarray<T>& arr)
-{
-  auto s = arr.shape();
-  std::vector<hsize_t> dims {s.cbegin(), s.cend()};
-  write_dataset(obj_id, dims.size(), dims.data(), name, H5TypeMap<T>::type_id,
-                arr.data(), false);
-}
-
-template<typename T, std::size_t N> inline void
-write_dataset(hid_t obj_id, const char* name, const xt::xtensor<T, N>& arr)
+// Template for xarray, xtensor, etc.
+template<typename D> inline void
+write_dataset(hid_t obj_id, const char* name, const xt::xcontainer<D>& arr)
 {
+  using T = typename D::value_type;
   auto s = arr.shape();
   std::vector<hsize_t> dims {s.cbegin(), s.cend()};
   write_dataset(obj_id, dims.size(), dims.data(), name, H5TypeMap<T>::type_id,

include/openmc/tallies/tally.h

@@ -9,9 +9,11 @@ namespace openmc {
 
 extern "C" double total_weight;
 
-xt::xtensor<double, 3> tally_results(int idx);
+template <std::size_t N>
+using adaptor_type = xt::xtensor_adaptor<xt::xbuffer_adaptor<double*&, xt::no_ownership>, N>;
 
-xt::xtensor<double, 2> global_tallies();
+adaptor_type<2> global_tallies();
+adaptor_type<3> tally_results(int idx);
 
 #ifdef OPENMC_MPI
 //! Collect all tally results onto master process

openmc/capi/tally.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping
-from ctypes import c_int, c_int32, c_double, c_char_p, c_bool, POINTER
+from ctypes import c_int, c_int32, c_size_t, c_double, c_char_p, c_bool, POINTER
 from weakref import WeakValueDictionary
 
 import numpy as np
@@ -60,7 +60,7 @@ _dll.openmc_tally_reset.argtypes = [c_int32]
 _dll.openmc_tally_reset.restype = c_int
 _dll.openmc_tally_reset.errcheck = _error_handler
 _dll.openmc_tally_results.argtypes = [
-    c_int32, POINTER(POINTER(c_double)), POINTER(c_int*3)]
+    c_int32, POINTER(POINTER(c_double)), POINTER(c_size_t*3)]
 _dll.openmc_tally_results.restype = c_int
 _dll.openmc_tally_results.errcheck = _error_handler
 _dll.openmc_tally_set_active.argtypes = [c_int32, c_bool]
@@ -296,7 +296,7 @@ class Tally(_FortranObjectWithID):
     @property
     def results(self):
         data = POINTER(c_double)()
-        shape = (c_int*3)()
+        shape = (c_size_t*3)()
         _dll.openmc_tally_results(self._index, data, shape)
         return as_array(data, tuple(shape))
 

src/simulation.cpp

@@ -3,10 +3,7 @@
 #include "openmc/capi.h"
 #include "openmc/message_passing.h"
 #include "openmc/settings.h"
-
-#ifdef OPENMC_MPI
-#include <mpi.h>
-#endif
+#include "openmc/tallies/tally.h"
 
 #include <algorithm>
 
@@ -106,29 +103,25 @@ void calculate_work()
 
 #ifdef OPENMC_MPI
 void broadcast_results() {
-    // Broadcast tally results so that each process has access to results
-  double* buffer;
-  if (n_tallies > 0) {
-    for (int i=1; i <= n_tallies; ++i) {
-      // Create a new datatype that consists of all values for a given filter
-      // bin and then use that to broadcast. This is done to minimize the
-      // chance of the 'count' argument of MPI_BCAST exceeding 2**31
-      int shape[3];
-      openmc_tally_results(i, &buffer, shape);
+  // Broadcast tally results so that each process has access to results
+  for (int i = 1; i <= n_tallies; ++i) {
+    // Create a new datatype that consists of all values for a given filter
+    // bin and then use that to broadcast. This is done to minimize the
+    // chance of the 'count' argument of MPI_BCAST exceeding 2**31
+    auto results = tally_results(i);
 
-      int n = shape[0];
-      int count_per_filter = shape[1] * shape[2];
-      MPI_Datatype result_block;
-      MPI_Type_contiguous(count_per_filter, MPI_DOUBLE, &result_block);
-      MPI_Type_commit(&result_block);
-      MPI_Bcast(buffer, n, result_block, 0, mpi::intracomm);
-      MPI_Type_free(&result_block);
-    }
+    auto shape = results.shape();
+    int count_per_filter = shape[1] * shape[2];
+    MPI_Datatype result_block;
+    MPI_Type_contiguous(count_per_filter, MPI_DOUBLE, &result_block);
+    MPI_Type_commit(&result_block);
+    MPI_Bcast(results.data(), shape[0], result_block, 0, mpi::intracomm);
+    MPI_Type_free(&result_block);
   }
 
   // Also broadcast global tally results
-  openmc_global_tallies(&buffer);
-  MPI_Bcast(buffer, 4, MPI_DOUBLE, 0, mpi::intracomm);
+  auto gt = global_tallies();
+  MPI_Bcast(gt.data(), gt.size(), MPI_DOUBLE, 0, mpi::intracomm);
 
   // These guys are needed so that non-master processes can calculate the
   // combined estimate of k-effective

src/tallies/tally.cpp

@@ -8,31 +8,33 @@
 #include "xtensor/xbuilder.hpp" // for empty_like
 #include "xtensor/xview.hpp"
 
+#include <array>
 #include <cstddef>
 
 namespace openmc {
 
-xt::xtensor<double, 2> global_tallies()
+adaptor_type<2> global_tallies()
 {
   // Get pointer to global tallies
   double* buffer;
   openmc_global_tallies(&buffer);
 
   // Adapt into xtensor
-  std::vector<int> shape {N_GLOBAL_TALLIES, 3};
-  int size {3*N_GLOBAL_TALLIES};
+  std::array<size_t, 2> shape = {N_GLOBAL_TALLIES, 3};
+  std::size_t size {3*N_GLOBAL_TALLIES};
+
   return xt::adapt(buffer, size, xt::no_ownership(), shape);
 }
 
-xt::xtensor<double, 3> tally_results(int idx)
+adaptor_type<3> tally_results(int idx)
 {
   // Get pointer to tally results
   double* results;
-  std::vector<int> shape(3);
+  std::array<std::size_t, 3> shape;
   openmc_tally_results(idx, &results, shape.data());
 
   // Adapt array into xtensor with no ownership
-  int size {shape[0] * shape[1] * shape[2]};
+  std::size_t size {shape[0] * shape[1] * shape[2]};
   return xt::adapt(results, size, xt::no_ownership(), shape);
 }
 
@@ -59,7 +61,7 @@ void reduce_tally_results()
 
     // Transfer values on master and reset on other ranks
     if (mpi::master) {
-      values_view = mpi::master;
+      values_view = values_reduced;
     } else {
       values_view = 0.0;
     }

src/tallies/tally_header.F90

@@ -681,7 +681,7 @@ contains
     ! allows a user to obtain in-memory tally results from Python directly.
     integer(C_INT32_T), intent(in), value :: index
     type(C_PTR),        intent(out) :: ptr
-    integer(C_INT),     intent(out) :: shape_(3)
+    integer(C_SIZE_T),  intent(out) :: shape_(3)
     integer(C_INT) :: err
 
     if (index >= 1 .and. index <= size(tallies)) then