add support for different data types in arrays, also on GPU

CarpetX/src/driver.cxx

@@ -1,3 +1,4 @@
+#pragma GCC optimize("O0")
 #include "driver.hxx"
 
 #include "boundaries.hxx"
@@ -1783,7 +1784,6 @@ void SetupGlobals() {
     if (group.grouptype == CCTK_GF)
       continue;
     assert(group.grouptype == CCTK_ARRAY || group.grouptype == CCTK_SCALAR);
-    assert(group.vartype == CCTK_VARIABLE_REAL);
     assert(group.disttype == CCTK_DISTRIB_CONSTANT);
     assert(group.dim >= 0);
     assert(group.dim <= dim);
@@ -1834,9 +1834,8 @@ void SetupGlobals() {
     arraygroupdata.data.resize(group.numtimelevels);
     arraygroupdata.valid.resize(group.numtimelevels);
     for (int tl = 0; tl < int(arraygroupdata.data.size()); ++tl) {
-      // TODO: Allocate in managed memory
-      arraygroupdata.data.at(tl).resize(arraygroupdata.numvars *
-                                        arraygroupdata.array_size);
+      arraygroupdata.data.at(tl).alloc(
+          group.vartype, arraygroupdata.numvars * arraygroupdata.array_size);
       why_valid_t why([]() { return "SetupGlobals"; });
       arraygroupdata.valid.at(tl).resize(arraygroupdata.numvars, why);
       for (int vi = 0; vi < arraygroupdata.numvars; ++vi) {
@@ -2074,7 +2073,7 @@ void CactusAmrCore::RemakeLevel(const int level, const amrex::Real time,
         check_valid_gf(active_levels, gi, vi, tl, nan_handling,
                        []() { return "RemakeLevel before prolongation"; });
       } // for vi
-    } // for tl
+    }   // for tl
 
   } // for gi
 
@@ -2291,6 +2290,28 @@ YAML::Emitter &operator<<(YAML::Emitter &yaml, const amrex::AmrCore &amrcore) {
 } // namespace amrex
 namespace CarpetX {
 
+std::ostream &
+operator<<(std::ostream &os,
+           const GHExt::GlobalData::AnyTypeVector::AnyTypeScalarRef &scalar) {
+  const char sep = '\t';
+  switch (scalar._vect.type()) {
+  case CCTK_VARIABLE_REAL:
+    os << *(CCTK_REAL *)scalar._vect.data_at(scalar._idx);
+    break;
+  case CCTK_VARIABLE_INT:
+    os << *(CCTK_INT *)scalar._vect.data_at(scalar._idx);
+    break;
+  case CCTK_VARIABLE_COMPLEX: {
+    CCTK_COMPLEX value = *(CCTK_COMPLEX *)scalar._vect.data_at(scalar._idx);
+    os << value.real() << sep << value.imag();
+  } break;
+  default:
+    assert(0 && "Unexpected variable type");
+    break;
+  }
+  return os;
+}
+
 YAML::Emitter &operator<<(YAML::Emitter &yaml,
                           const GHExt::CommonGroupData &commongroupdata) {
   yaml << YAML::LocalTag("commongroupdata-1.0.0");
@@ -2312,6 +2333,47 @@ YAML::Emitter &operator<<(YAML::Emitter &yaml,
   return yaml;
 }
 
+YAML::Emitter &operator<<(YAML::Emitter &yaml, const CCTK_COMPLEX &cval) {
+  yaml << YAML::Flow << YAML::BeginSeq << cval.real() << cval.imag()
+       << YAML::EndSeq;
+  return yaml;
+};
+
+YAML::Emitter &
+operator<<(YAML::Emitter &yaml,
+           const GHExt::GlobalData::AnyTypeVector::AnyTypeScalarRef
+               &anytypescalarref) {
+  switch (anytypescalarref._vect.type()) {
+  case CCTK_VARIABLE_COMPLEX:
+    yaml << *(const CCTK_COMPLEX *)anytypescalarref._vect.data_at(
+        anytypescalarref._idx);
+    break;
+  case CCTK_VARIABLE_REAL:
+    yaml << *(const CCTK_REAL *)anytypescalarref._vect.data_at(
+        anytypescalarref._idx);
+    break;
+  case CCTK_VARIABLE_INT:
+    yaml << *(const CCTK_INT *)anytypescalarref._vect.data_at(
+        anytypescalarref._idx);
+    break;
+  default:
+    // missed to implement a type
+    CCTK_VERROR("Cannot handle type %d", anytypescalarref._vect.type());
+    break;
+  }
+  return yaml;
+}
+
+YAML::Emitter &
+operator<<(YAML::Emitter &yaml,
+           const GHExt::GlobalData::AnyTypeVector &anytypevector) {
+  yaml << YAML::BeginSeq;
+  for (size_t i = 0; i < anytypevector.size(); ++i)
+    yaml << anytypevector[i];
+  yaml << YAML::EndSeq;
+  return yaml;
+}
+
 YAML::Emitter &
 operator<<(YAML::Emitter &yaml,
            const GHExt::GlobalData::ArrayGroupData &arraygroupdata) {

CarpetX/src/driver.hxx

@@ -152,13 +152,145 @@ struct GHExt {
   struct GlobalData {
     // all data that exists on all levels
 
+    class AnyTypeVector {
+
+    public:
+      // access to a single element of a AnyTypeVector
+      class AnyTypeScalarRef {
+      public:
+        AnyTypeScalarRef() = delete;
+        AnyTypeScalarRef(const AnyTypeVector &vect_, size_t idx_)
+            : _vect(vect_), _idx(idx_) {}
+
+      private:
+        const AnyTypeVector &_vect;
+        const size_t _idx;
+
+        friend YAML::Emitter &
+        operator<<(YAML::Emitter &yaml,
+                   const AnyTypeScalarRef &anytypescalarref);
+        friend std::ostream &operator<<(std::ostream &os,
+                                        const AnyTypeScalarRef &scalar);
+      };
+
+      AnyTypeVector() : _type(-1), _typesize(-1), _count(0), _data(nullptr){};
+      AnyTypeVector(int type_, size_t count_)
+          : _type(-1), _typesize(-1), _count(0), _data(nullptr) {
+        alloc(type_, count_);
+        assert(_type == type_);
+        assert(_typesize != -1);
+        assert(_count == count_);
+        assert(_data != nullptr);
+      };
+      // Noncopyable for now
+      AnyTypeVector(const AnyTypeVector &) = delete;
+      AnyTypeVector &operator=(const AnyTypeVector &) = delete;
+      AnyTypeVector &operator=(AnyTypeVector &&other) {
+        swap(other);
+        return *this;
+      }
+      AnyTypeVector(AnyTypeVector &&other)
+          : _type(other._type), _typesize(other._typesize),
+            _count(other._count), _data(other._data) {
+        other._type = -1;
+        other._typesize = -1;
+        other._count = 0;
+        other._data = nullptr;
+      }
+      void swap(AnyTypeVector &other) {
+        std::swap(this->_type, other._type);
+        std::swap(this->_typesize, other._typesize);
+        std::swap(this->_count, other._count);
+        std::swap(this->_data, other._data);
+      }
+
+      ~AnyTypeVector() {
+        if (_data != nullptr) {
+          assert(_type != -1);
+          assert(_typesize != -1);
+          amrex::The_Arena()->free(_data);
+          _type = -1;
+          _typesize = -1;
+          _count = 0;
+          _data = nullptr;
+        }
+        assert(_type == -1);
+        assert(_typesize == -1);
+        assert(_count == 0);
+        assert(_data == nullptr);
+      };
+
+      void alloc(int type_, size_t count_) {
+        assert(type_ == CCTK_VARIABLE_INT || type_ == CCTK_VARIABLE_REAL ||
+               type_ == CCTK_VARIABLE_COMPLEX);
+
+        assert(_type == -1);
+        assert(_typesize == -1);
+        assert(_count == 0);
+        assert(_data == nullptr);
+
+        _type = type_;
+        _typesize = CCTK_VarTypeSize(_type);
+        assert(_typesize > 0);
+        _count = count_;
+        _data = amrex::The_Arena()->alloc(_typesize * _count);
+      }
+
+      void free() {
+        assert(_type != -1);
+        assert(_typesize != -1);
+        assert(_data != nullptr);
+        amrex::The_Arena()->free(_data);
+        _type = -1;
+        _typesize = -1;
+        _count = 0;
+        _data = nullptr;
+      }
+
+      int type() const { return _type; };
+      int typesize() const { return _typesize; };
+
+      const void *data_at(size_t i) const {
+#ifdef CCTK_DEBUG
+        if (i >= _count) {
+          CCTK_VERROR("invalid index %zd exceeds %zd", i, _count);
+        }
+#endif
+        assert(i < _count);
+        return (char *)_data + i * _typesize;
+      };
+
+      void *data_at(size_t i) {
+#ifdef CCTK_DEBUG
+        if (i >= _count) {
+          CCTK_VERROR("invalid index %zu exceeds %zu", i, _count);
+        }
+#endif
+        assert(i < _count);
+        return (char *)_data + i * _typesize;
+      };
+
+      AnyTypeScalarRef operator[](size_t idx) const {
+        return AnyTypeScalarRef(*this, idx);
+      }
+
+      size_t size() const { return _count; };
+
+      friend YAML::Emitter &operator<<(YAML::Emitter &yaml,
+                                       const AnyTypeVector &anytypevector);
+
+    private:
+      int _type, _typesize;
+      size_t _count;
+      void *_data;
+    };
+
     // For subcycling in time, there really should be one copy of each
     // integrated grid scalar per level. We don't do that yet; instead,
     // we assume that grid scalars only hold "analysis" data.
 
     struct ArrayGroupData : public CommonGroupData {
-      std::vector<std::vector<CCTK_REAL> >
-          data; // [time level][var index + grid point index]
+      vector<AnyTypeVector> data; // [time level][var index + grid point index]
       int array_size;
       int dimension;
       int activetimelevels;