Add implementation for NLC functionals (VV10)

include/gauxc/reduction_driver.hpp

@@ -13,8 +13,11 @@
 #include <memory>
 #include <gauxc/gauxc_config.hpp>
 #include <gauxc/runtime_environment.hpp>
+#include <cstring>
 #include <typeindex>
 #include <any>
+#include <cstddef>
+#include <vector>
 
 namespace GauXC {
 
@@ -52,11 +55,23 @@ class ReductionDriver {
     allreduce_inplace_typeerased( data, size, op, std::type_index(typeid(T)), optional_args );
   }
 
+  template <typename T>
+  inline std::vector<T> allgather_v( const T* src, size_t size, std::any optional_args = std::any() ) const {
+    std::vector<std::byte> bytes;
+    allgather_v_typeerased( src, size, bytes, std::type_index(typeid(T)), optional_args );
+    std::vector<T> dest( bytes.size() / sizeof(T) );
+    std::memcpy( dest.data(), bytes.data(), bytes.size() );
+    return dest;
+  }
+
   void allreduce_typeerased( const void*, void*, size_t, ReductionOp, std::type_index, std::any );
   void allreduce_inplace_typeerased( void*, size_t, ReductionOp, std::type_index, std::any );
+  void allgather_v_typeerased( const void*, size_t, std::vector<std::byte>&, std::type_index, std::any ) const;
 
   bool takes_host_memory() const;
   bool takes_device_memory() const;
+  int comm_rank() const;
+  int comm_size() const;
 
 };
 

include/gauxc/xc_integrator.hpp

@@ -80,6 +80,18 @@ class XCIntegrator {
   exx_type      eval_exx     ( const MatrixType&, 
                                const IntegratorSettingsEXX& = IntegratorSettingsEXX{} );
 
+  value_type eval_nlc( const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+  value_type eval_nlc( const MatrixType&, const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+
+  exc_vxc_type_rks eval_nlc_vnlc( const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+  exc_vxc_type_uks eval_nlc_vnlc( const MatrixType&, const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+
+  exc_grad_type eval_nlc_grad( const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+  exc_grad_type eval_nlc_grad( const MatrixType&, const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+
+  fxc_contraction_type_rks eval_nlc_fnlc_contraction( const MatrixType&, const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+  fxc_contraction_type_uks eval_nlc_fnlc_contraction( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&, const IntegratorSettingsNLC& = IntegratorSettingsNLC{} );
+
   fxc_contraction_type_rks  eval_fxc_contraction ( const MatrixType&, const MatrixType&,
                                   const IntegratorSettingsXC& = IntegratorSettingsXC{} );
   fxc_contraction_type_uks  eval_fxc_contraction ( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&,

include/gauxc/xc_integrator/impl.hpp

@@ -100,6 +100,63 @@ typename XCIntegrator<MatrixType>::exx_type
   return pimpl_->eval_exx(P,settings);
 };
 
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::value_type
+  XCIntegrator<MatrixType>::eval_nlc( const MatrixType& P, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc(P, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::value_type
+  XCIntegrator<MatrixType>::eval_nlc( const MatrixType& Ps, const MatrixType& Pz, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc(Ps, Pz, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_vxc_type_rks
+  XCIntegrator<MatrixType>::eval_nlc_vnlc( const MatrixType& P, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_vnlc(P, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_vxc_type_uks
+  XCIntegrator<MatrixType>::eval_nlc_vnlc( const MatrixType& Ps, const MatrixType& Pz, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_vnlc(Ps, Pz, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_grad_type
+  XCIntegrator<MatrixType>::eval_nlc_grad( const MatrixType& P, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_grad(P, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_grad_type
+  XCIntegrator<MatrixType>::eval_nlc_grad( const MatrixType& Ps, const MatrixType& Pz, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_grad(Ps, Pz, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::fxc_contraction_type_rks
+  XCIntegrator<MatrixType>::eval_nlc_fnlc_contraction( const MatrixType& P, const MatrixType& tP, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_fnlc_contraction(P, tP, settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::fxc_contraction_type_uks
+  XCIntegrator<MatrixType>::eval_nlc_fnlc_contraction( const MatrixType& Ps, const MatrixType& Pz,
+                           const MatrixType& tPs, const MatrixType& tPz, const IntegratorSettingsNLC& settings ) {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_nlc_fnlc_contraction(Ps, Pz, tPs, tPz, settings);
+};
+
 template <typename MatrixType>
 typename XCIntegrator<MatrixType>::fxc_contraction_type_rks
   XCIntegrator<MatrixType>::eval_fxc_contraction( const MatrixType& P, const MatrixType& tP, 

include/gauxc/xc_integrator/replicated/replicated_xc_integrator_impl.hpp

@@ -176,7 +176,6 @@ class ReplicatedXCIntegratorImpl {
                       value_type* FXCs, int64_t ldfxcs,
                       value_type* FXCz, int64_t ldfxcz,
                       const IntegratorSettingsXC& ks_settings );
-
   void eval_dd_psi( int64_t m, int64_t n, const value_type* P,
                      int64_t ldp, unsigned max_Ylm, 
                      value_type* ddPsi, int64_t ldPsi );

include/gauxc/xc_integrator/xc_integrator_impl.hpp

@@ -151,6 +151,53 @@ class XCIntegratorImpl {
     return eval_exx_(P,settings);
   }
 
+  value_type eval_nlc( const MatrixType& P, const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_(P, nlc_settings);
+  }
+
+  value_type eval_nlc( const MatrixType& Ps, const MatrixType& Pz,
+                           const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_(Ps, Pz, nlc_settings);
+  }
+
+  exc_vxc_type_rks eval_nlc_vnlc( const MatrixType& P,
+                                     const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_vxc_(P, nlc_settings);
+  }
+
+  exc_vxc_type_uks eval_nlc_vnlc( const MatrixType& Ps, const MatrixType& Pz,
+                                     const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_vxc_(Ps, Pz, nlc_settings);
+  }
+
+  exc_grad_type eval_nlc_grad( const MatrixType& P,
+                                   const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_grad_(P, nlc_settings);
+  }
+
+  exc_grad_type eval_nlc_grad( const MatrixType& Ps, const MatrixType& Pz,
+                                   const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_exc_grad_(Ps, Pz, nlc_settings);
+  }
+
+  fxc_contraction_type_rks eval_nlc_fnlc_contraction( const MatrixType& P,
+    const MatrixType& tP, const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_fxc_contraction_(P, tP, nlc_settings);
+  }
+
+  fxc_contraction_type_uks eval_nlc_fnlc_contraction( const MatrixType& Ps, const MatrixType& Pz,
+    const MatrixType& tPs, const MatrixType& tPz, const IntegratorSettingsNLC& settings ) {
+    detail::IntegratorSettingsNLCInternal nlc_settings(settings);
+    return eval_fxc_contraction_(Ps, Pz, tPs, tPz, nlc_settings);
+  }
+
 
   /** Integrate FXC contraction for RKS
    * 

include/gauxc/xc_integrator_settings.hpp

@@ -29,4 +29,25 @@ struct IntegratorSettingsEXC_GRAD : public IntegratorSettingsKS {
   bool include_weight_derivatives= true; // whether to include grid weight contribution and employ translational invariance, or just use Hellmann-Feynman gradient
 };
 
+enum class NLCMathMode {
+  NativeFP64,
+  FloatPair
+};
+
+struct IntegratorSettingsNLC : public IntegratorSettingsEXC_GRAD {
+  double vv10_b      = 6.3;
+  double vv10_c      = 0.0093;
+  double vv10_tol    = 1e-8;
+  NLCMathMode math_mode = NLCMathMode::NativeFP64;
+};
+
+namespace detail {
+
+struct IntegratorSettingsNLCInternal : public IntegratorSettingsNLC {
+  IntegratorSettingsNLCInternal() = default;
+  explicit IntegratorSettingsNLCInternal( const IntegratorSettingsNLC& settings ) :
+    IntegratorSettingsNLC(settings) { }
+};
+
+}
 }

src/reduction_driver/host/basic_mpi_reduction_driver.cxx

@@ -10,11 +10,14 @@
  * See LICENSE.txt for details
  */
 #include "basic_mpi_reduction_driver.hpp"
+#include <gauxc/exceptions.hpp>
 #include <cstring>
 #include <memory>
 #include <map>
 #include <iostream>
 #include <cstddef>
+#include <limits>
+#include <vector>
 
 namespace GauXC {
 
@@ -23,7 +26,8 @@ MPI_Datatype get_mpi_datatype( std::type_index idx ) {
 
   static std::map<std::type_index, MPI_Datatype> map {
     {std::type_index(typeid(double)), MPI_DOUBLE},
-    {std::type_index(typeid(float)),  MPI_FLOAT}
+    {std::type_index(typeid(float)),  MPI_FLOAT},
+    {std::type_index(typeid(unsigned long long)), MPI_UNSIGNED_LONG_LONG}
   };
 
   return map.at(idx);
@@ -45,7 +49,8 @@ size_t get_dtype_size( std::type_index idx ) {
 
   static std::map<std::type_index, size_t> map {
     {std::type_index(typeid(double)), sizeof(double)}, 
-    {std::type_index(typeid(float)),  sizeof(float)}
+    {std::type_index(typeid(float)),  sizeof(float)},
+    {std::type_index(typeid(unsigned long long)), sizeof(unsigned long long)}
   };
 
   return map.at(idx);
@@ -108,6 +113,42 @@ void BasicMPIReductionDriver::allreduce_inplace_typeerased( void* data, size_t s
   }
 }
 
+void BasicMPIReductionDriver::allgather_v_typeerased( const void* src, size_t size,
+  std::vector<std::byte>& dest, std::type_index idx, std::any optional_args ) {
+
+  if( optional_args.has_value() )
+    std::cout << "** Warning: Optional Args Are Not Used in BasiMPIReductionDriver::allgather_v" << std::endl;
+
+  const auto dtype_size = get_dtype_size(idx);
+  const int world_size = runtime_.comm_size();
+  if( world_size == 1 ) {
+    dest.resize( size * dtype_size );
+    std::memcpy( dest.data(), src, dest.size() );
+    return;
+  }
+
+#ifdef GAUXC_HAS_MPI
+  std::vector<unsigned long long> sizes(world_size);
+  const auto local_size = static_cast<unsigned long long>(size);
+  MPI_Allgather( &local_size, 1, MPI_UNSIGNED_LONG_LONG, sizes.data(), 1,
+                 MPI_UNSIGNED_LONG_LONG, runtime_.comm() );
+
+  std::vector<int> counts(world_size), displs(world_size);
+  size_t total_size = 0;
+  for( int i = 0; i < world_size; ++i ) {
+    if( sizes[i] > static_cast<unsigned long long>(std::numeric_limits<int>::max()) )
+      GAUXC_GENERIC_EXCEPTION("allgather_v message count exceeds MPI int range");
+    counts[i] = static_cast<int>( sizes[i] );
+    displs[i] = static_cast<int>( total_size );
+    total_size += static_cast<size_t>( sizes[i] );
+  }
+
+  dest.resize( total_size * dtype_size );
+  MPI_Allgatherv( src, static_cast<int>(size), get_mpi_datatype(idx), dest.data(),
+                  counts.data(), displs.data(), get_mpi_datatype(idx), runtime_.comm() );
+#endif
+}
+
 std::unique_ptr<detail::ReductionDriverImpl> BasicMPIReductionDriver::clone() {
   return std::make_unique<BasicMPIReductionDriver>(*this);
 }

src/reduction_driver/host/basic_mpi_reduction_driver.hpp

@@ -22,6 +22,7 @@ struct BasicMPIReductionDriver : public HostReductionDriver {
 
   void allreduce_typeerased( const void*, void*, size_t, ReductionOp, std::type_index, std::any ) override;
   void allreduce_inplace_typeerased( void*, size_t, ReductionOp, std::type_index, std::any ) override;
+  void allgather_v_typeerased( const void*, size_t, std::vector<std::byte>&, std::type_index, std::any ) override;
 
   std::unique_ptr<detail::ReductionDriverImpl> clone() override;
 

src/reduction_driver/reduction_driver.cxx

@@ -38,6 +38,11 @@ void ReductionDriver::allreduce_inplace_typeerased( void* data, size_t size, Red
   pimpl_->allreduce_inplace_typeerased(data, size, op, idx, optional_args);
 }
 
+void ReductionDriver::allgather_v_typeerased( const void* src, size_t size, std::vector<std::byte>& dest, std::type_index idx, std::any optional_args ) const {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  pimpl_->allgather_v_typeerased(src, size, dest, idx, optional_args);
+}
+
 bool ReductionDriver::takes_host_memory() const {
   if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
   return pimpl_->takes_host_memory();
@@ -47,6 +52,16 @@ bool ReductionDriver::takes_device_memory() const {
   return pimpl_->takes_device_memory();
 }
 
+int ReductionDriver::comm_rank() const {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->comm_rank();
+}
+
+int ReductionDriver::comm_size() const {
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->comm_size();
+}
+
 
 
 

src/reduction_driver/reduction_driver_factory.cxx

@@ -23,6 +23,13 @@
 
 namespace GauXC {
 
+ReductionDriver ReductionDriverFactory::get_instance(
+  const RuntimeEnvironment& rt, std::string kernel_name ) {
+
+  auto driver = get_shared_instance(rt, std::move(kernel_name));
+  return *driver;
+}
+
 std::shared_ptr<ReductionDriver> ReductionDriverFactory::get_shared_instance(
   const RuntimeEnvironment& rt, std::string kernel_name ) {
 

src/reduction_driver/reduction_driver_impl.cxx

@@ -10,6 +10,7 @@
  * See LICENSE.txt for details
  */
 #include "reduction_driver_impl.hpp"
+#include <gauxc/exceptions.hpp>
 
 namespace GauXC::detail {
 
@@ -20,4 +21,12 @@ ReductionDriverImpl::ReductionDriverImpl( const RuntimeEnvironment& rt )
 ReductionDriverImpl::~ReductionDriverImpl() noexcept = default;
 ReductionDriverImpl::ReductionDriverImpl(const ReductionDriverImpl& ) = default;
 
+void ReductionDriverImpl::allgather_v_typeerased( const void*, size_t,
+  std::vector<std::byte>&, std::type_index, std::any ) {
+  GAUXC_GENERIC_EXCEPTION("Variable-size allgather is not supported by this ReductionDriver");
+}
+
+int ReductionDriverImpl::comm_rank() const { return runtime_.comm_rank(); }
+int ReductionDriverImpl::comm_size() const { return runtime_.comm_size(); }
+
 }

src/reduction_driver/reduction_driver_impl.hpp

@@ -11,6 +11,8 @@
  */
 #pragma once
 #include <gauxc/reduction_driver.hpp>
+#include <cstddef>
+#include <vector>
 
 namespace GauXC  {
 namespace detail {
@@ -31,9 +33,12 @@ class ReductionDriverImpl {
 
   virtual void allreduce_typeerased( const void*, void*, size_t, ReductionOp, std::type_index, std::any ) = 0;
   virtual void allreduce_inplace_typeerased( void*, size_t, ReductionOp, std::type_index, std::any ) = 0;
+  virtual void allgather_v_typeerased( const void*, size_t, std::vector<std::byte>&, std::type_index, std::any );
 
   virtual bool takes_host_memory() const = 0;
   virtual bool takes_device_memory() const = 0;
+  int comm_rank() const;
+  int comm_size() const;
 
   virtual std::unique_ptr<ReductionDriverImpl> clone() = 0;
 };

src/xc_integrator/integrator_util/CMakeLists.txt

@@ -9,4 +9,4 @@
 #
 # See LICENSE.txt for details
 #
-target_sources( gauxc PRIVATE integrator_common.cxx integral_bounds.cxx exx_screening.cxx spherical_harmonics.cxx )
+target_sources( gauxc PRIVATE integrator_common.cxx integral_bounds.cxx exx_screening.cxx spherical_harmonics.cxx vv10.cxx )