From b0330a36f72a6073f99a4e315a906fc3c2232e87 Mon Sep 17 00:00:00 2001
From: Russell Standish <hpcoder@hpcoders.com.au>
Date: Sat, 7 Dec 2024 09:29:55 +1100
Subject: [PATCH] WIP for getting model components in the right memory space.
 Still have an issue with Ouroboros-SYCL not being able to allocate more than
 16 bytes at a time.

---
 classdesc              |   2 +-
 graphcode              |   2 +-
 include/ecolab.h       | 110 ++++++++++++++++++++++-------------------
 models/ecolab_model.cc |  14 ++++--
 models/ecolab_model.h  |   6 +--
 5 files changed, 74 insertions(+), 60 deletions(-)

diff --git a/classdesc b/classdesc
index 392d56c..c2be4af 160000
--- a/classdesc
+++ b/classdesc
@@ -1 +1 @@
-Subproject commit 392d56c6e560c9d5ed49287701edb1a100b75687
+Subproject commit c2be4afd64d9652ac76b67dd33be383085b79203
diff --git a/graphcode b/graphcode
index 33a2f47..db52ce5 160000
--- a/graphcode
+++ b/graphcode
@@ -1 +1 @@
-Subproject commit 33a2f47f38981960401756b92929b619760ee99f
+Subproject commit db52ce5a8ee572ec964b876021821ba01a4861a5
diff --git a/include/ecolab.h b/include/ecolab.h
index fffbb1e..8412a50 100644
--- a/include/ecolab.h
+++ b/include/ecolab.h
@@ -133,36 +133,40 @@ namespace ecolab
     using MemAllocator=Ouro::MultiOuroPQ;
     MemAllocator* memAlloc=nullptr;
     const sycl::stream* out=nullptr;
-    template <class T> class Allocator
+    template <class T> class CellAllocator
     {
     public:
       Ouro::SyclDesc<1>*const* desc=nullptr;
       MemAllocator*const* memAlloc;
+      const sycl::stream* out=nullptr;
       template <class U> friend class Allocator;
-      Allocator()=default;
-      Allocator(Ouro::SyclDesc<1>* const& desc, MemAllocator*const& memAlloc):
-        desc(&desc), memAlloc(&memAlloc) {
+      CellAllocator()=default;
+      CellAllocator(Ouro::SyclDesc<1>* const& desc, MemAllocator*const& memAlloc, const sycl::stream* out):
+        desc(&desc), memAlloc(&memAlloc), out(out) {
       }
-      template <class U> Allocator(const Allocator<U>& x):
+      template <class U> CellAllocator(const CellAllocator<U>& x):
         desc(x.desc) {}
       T* allocate(size_t sz) {
-        if (memAlloc && *memAlloc && desc && *desc) 
-          return reinterpret_cast<T*>((*memAlloc)->malloc(**desc,sz*sizeof(T)));
-        else
-          return nullptr; // TODO raise an error??
+        if (memAlloc && *memAlloc && desc && *desc)  {
+          auto r=reinterpret_cast<T*>((*memAlloc)->malloc(**desc,sz*sizeof(T)));
+          if (out) (*out)<<"allocated "<<sz*sizeof(T)<<" bytes, "<<r<<sycl::endl;
+          return r;
+        }
+        if (out) (*out)<<"failed allocation "<<sz*sizeof(T)<<sycl::endl;
+        return nullptr; // TODO raise an error??
       }
       void deallocate(T* p,size_t) {
         if (memAlloc && *memAlloc && desc && *desc)
           (*memAlloc)->free(**desc,p);
       }
-      bool operator==(const Allocator& x) const {return desc==x.desc && memAlloc==x.memAlloc;}
+      bool operator==(const CellAllocator& x) const {return desc==x.desc && memAlloc==x.memAlloc;}
     };
-    template <class T> Allocator<T> allocator() const {
-      return Allocator<T>(desc,memAlloc);
+    template <class T> CellAllocator<T> allocator() const {
+      return CellAllocator<T>(desc,memAlloc,out);
     }
 #else
-    template <class T> using Allocator=std::allocator<T>;
-    template <class T> Allocator<T> allocator() const {return Allocator<T>();}
+    template <class T> using CellAllocator=std::allocator<T>;
+    template <class T> CellAllocator<T> allocator() const {return CellAllocator<T>();}
 #endif
   };
 
@@ -187,31 +191,15 @@ namespace ecolab
 #endif      
   };
 
-  // This class exists, because we need a default constructor, and
-  // usm_allocator needs to be initialised with a device context
-#ifdef SYCL_LANGUAGE_VERSION
-  template <class T>
-  struct CellAllocator: public sycl::usm_allocator<T,sycl::usm::alloc::shared>
-  {
-    CellAllocator(): sycl::usm_allocator<T,sycl::usm::alloc::shared>(syclQ()) {}
-    template<class U> constexpr CellAllocator(const CellAllocator<U>& x) noexcept:
-      sycl::usm_allocator<T,sycl::usm::alloc::shared>(x) {}
-    template<class U> struct rebind {using other=CellAllocator<U>;};
-  };
-#else
-  template <class T> struct CellAllocator: std::allocator<T>
-  {
-    CellAllocator()=default;
-    template<class U> constexpr CellAllocator(const CellAllocator<U>& x) noexcept:
-      std::allocator<T>(x) {}
-    template<class U> struct rebind {using other=CellAllocator<U>;};
-  };
-#endif
-  
-  
   template <class Cell> struct EcolabGraph:
-    public Exclude<SyclGraphBase>, public graphcode::Graph<Cell, CellAllocator>
+    public Exclude<SyclGraphBase>, public graphcode::Graph<Cell>
   {
+#ifdef SYCL_LANGUAGE_VERSION
+    EcolabGraph(): graphcode::Graph<Cell>
+      (graphcode::Allocator<Cell>(syclQ(),sycl::usm::alloc::shared),
+       graphcode::Allocator<graphcode::ObjRef>(syclQ(),sycl::usm::alloc::shared),
+       graphcode::Allocator<graphcode::ObjectPtr<Cell>>(syclQ(),sycl::usm::alloc::shared)) {}
+#endif
     /// apply a functional to all local cells of this processor in parallel
     /// @param f 
     template <class F>
@@ -227,6 +215,7 @@ namespace ecolab
         h.parallel_for(sycl::nd_range<1>(range*workGroupSize, workGroupSize), [=,this](auto i) {
           auto idx=i.get_global_linear_id();
           if (idx<this->size()) {
+            out<<"idx="<<idx<<"(*this)[idx]"<<(*this)[idx].payload->get()<<" as "<<(*this)[idx]->template as<Cell>()<<sycl::endl;
             auto& cell=*(*this)[idx]->template as<Cell>();
             Ouro::SyclDesc<> desc(i,{});
             cell.desc=&desc;
@@ -251,15 +240,37 @@ namespace ecolab
 
 namespace classdesc
 {
-  template <class M> struct is_smart_ptr<ecolab::DeviceType<M>>: public true_type {}; 
+  template <class M> struct is_smart_ptr<ecolab::DeviceType<M>>: public true_type {};
+  template <class T> struct is_char: public false_type {};
+  template <> struct is_char<char>: public true_type {};
+  template <> struct is_char<wchar_t>: public true_type {};
+
+#ifdef SYCL_LANGUAGE_VERSION
+  /// classdesc support for stringifying usm_alloc
+  namespace   {
+    template <> EnumKey enum_keysData< ::sycl::usm::alloc >::keysData[]=
+      {
+        {"host",int(::sycl::usm::alloc::host)},
+        {"device",int(::sycl::usm::alloc::device)},
+        {"shared",int(::sycl::usm::alloc::shared)},
+        {"unknown",int(::sycl::usm::alloc::unknown)}
+      };
+    template <> EnumKeys< ::sycl::usm::alloc > enum_keysData< ::sycl::usm::alloc >::keys(enum_keysData< ::sycl::usm::alloc >::keysData,sizeof(enum_keysData< ::sycl::usm::alloc >::keysData)/sizeof(enum_keysData< ::sycl::usm::alloc >::keysData[0]));
+    template <> int enumKey< ::sycl::usm::alloc >(const std::string& x){return int(enum_keysData< ::sycl::usm::alloc >::keys(x));}
+    template <> std::string enumKey< ::sycl::usm::alloc >(int x){return enum_keysData< ::sycl::usm::alloc >::keys(x);}
+  }
+#endif
+
 }
 
-namespace classdesc_access
+namespace std
 {
   template <class T>
-  struct access_pack<ecolab::CellAllocator<T>>: public classdesc::NullDescriptor<classdesc::pack_t> {};
-  template <class T>
-  struct access_unpack<ecolab::CellAllocator<T>>: public classdesc::NullDescriptor<classdesc::pack_t> {};
+  typename classdesc::enable_if<
+    classdesc::Not<classdesc::is_char<typename std::remove_cv<T>::type>>,
+    std::ostream&>::T
+  operator<<(std::ostream& o, T* p)
+  {return o<<std::hex<<std::size_t(p);}
 }
 
 #ifdef MPI_SUPPORT
@@ -279,19 +290,14 @@ namespace ecolab
 #define CLASSDESC_unpack___ecolab__CellBase
 #define CLASSDESC_pack___ecolab__SyclGraphBase
 #define CLASSDESC_unpack___ecolab__SyclGraphBase
-#define CLASSDESC_pack___ecolab__CellBase__Allocator_T_
-#define CLASSDESC_unpack___ecolab__CellBase__Allocator_T_
-#define CLASSDESC_pack___ecolab__CellAllocator_T_
-#define CLASSDESC_unpack___ecolab__CellAllocator_T_
+#define CLASSDESC_pack___ecolab__CellBase__CellAllocator_T_
+#define CLASSDESC_unpack___ecolab__CellBase__CellAllocator_T_
 #define CLASSDESC_json_pack___ecolab__CellBase
 #define CLASSDESC_json_unpack___ecolab__CellBase
-#define CLASSDESC_json_pack___ecolab__CellBase__Allocator_T_
-#define CLASSDESC_json_unpack___ecolab__CellBase__Allocator_T_
-#define CLASSDESC_json_pack___ecolab__CellAllocator_T_
-#define CLASSDESC_json_unpack___ecolab__CellAllocator_T_
+#define CLASSDESC_json_pack___ecolab__CellBase__CellAllocator_T_
+#define CLASSDESC_json_unpack___ecolab__CellBase__CellAllocator_T_
 #define CLASSDESC_RESTProcess___ecolab__CellBase
-#define CLASSDESC_RESTProcess___ecolab__CellBase__Allocator_T_
-#define CLASSDESC_RESTProcess___ecolab__CellAllocator_T_
+#define CLASSDESC_RESTProcess___ecolab__CellBase__CellAllocator_T_
 
 #include "ecolab.cd"
 #endif  /* ECOLAB_H */
diff --git a/models/ecolab_model.cc b/models/ecolab_model.cc
index b841961..55b2155 100644
--- a/models/ecolab_model.cc
+++ b/models/ecolab_model.cc
@@ -479,6 +479,7 @@ void SpatialModel::setGrid(size_t nx, size_t ny)
         o->neighbours.push_back(makeId(i,j-1)); 
         o->neighbours.push_back(makeId(i,j+1)); 
       }
+  rebuildPtrLists();
 }
 
 void SpatialModel::generate(unsigned niter)
@@ -552,12 +553,19 @@ void SpatialModel::makeConsistent()
   forAll([=,this](EcolabCell& c) {
     if (nsp>c.density.size())
       {
-        array<int,EcolabCell::Allocator<int>> tmp(nsp,0,c.allocator<int>());
+        array<int,EcolabCell::CellAllocator<int>> tmp(nsp,0,c.allocator<int>());
         asg_v(tmp.data(),c.density.size(),c.density);
+        //(*c.out)<<"tmp.size:"<<tmp.size()<<" "<<tmp.data()<<sycl::endl;
         c.density.swap(tmp);
-        //(*c.out)<<"c.density.size:"<<c.density.size()<<" "<<c.density.data()<<sycl::endl;
+        //(*c.out)<<"c.density.size:"<<c.density.size()<<" "<<&c<<sycl::endl;
       }
   });
-  for (auto& i: *this) cout<<i.id()<<" "<<i->as<EcolabCell>()->density.size()<<endl;
+  for (auto& i: *this)
+    {
+      auto& c=*i->as<EcolabCell>();
+      //cout<<i.id()<<" "<<c.density.size()<<" "<<&c<<" usm type:"<<
+      //  sycl::get_pointer_type(&c,syclQ().get_context())<<" "<<sycl::get_pointer_type(i.payload,syclQ().get_context())<<endl;
+      
+    }
   ModelData::makeConsistent(nsp);
 }
diff --git a/models/ecolab_model.h b/models/ecolab_model.h
index 878743b..636ac43 100644
--- a/models/ecolab_model.h
+++ b/models/ecolab_model.h
@@ -68,7 +68,7 @@ template <class CellBase>
 struct EcolabPoint: public Exclude<CellBase>
 {
   Float salt;  /* random no. used for migration */
-  array<int,typename CellBase::template Allocator<int>> density{this->template allocator<int>()};
+  array<int,typename CellBase::template CellAllocator<int>> density{this->template allocator<int>()};
   void generate(unsigned niter, const ModelData&);
   void condense(const array<bool>& mask, size_t mask_true);
   array<int> mutate(const array<double>&);
@@ -82,8 +82,8 @@ struct EcolabPoint: public Exclude<CellBase>
 // for the panmictic model, we need to use std::allocator
 struct AllocatorBase
 {
-  template <class T> using Allocator=std::allocator<T>;
-  template <class T> Allocator<T> allocator() const {return Allocator<T>();}
+  template <class T> using CellAllocator=std::allocator<T>;
+  template <class T> CellAllocator<T> allocator() const {return CellAllocator<T>();}
 };
 
 struct PanmicticModel: public ModelData, public EcolabPoint<AllocatorBase>, public ecolab::Model<PanmicticModel>