Implement NEO-DFT on host

cmake/gauxc-dep-versions.cmake

@@ -11,7 +11,7 @@ set( GAUXC_CUTLASS_REPOSITORY https://github.com/NVIDIA/cutlass.git )
 set( GAUXC_CUTLASS_REVISION v2.10.0 )
 
 set( GAUXC_EXCHCXX_REPOSITORY https://github.com/wavefunction91/ExchCXX.git )
-set( GAUXC_EXCHCXX_REVISION   21a4700a826ec0beae1311a1d59677393bcb168f )
+set( GAUXC_EXCHCXX_REVISION   ef88e168249e38c526353fe7c48aab4bc697c99d )
 
 set( GAUXC_GAU2GRID_REPOSITORY https://github.com/dgasmith/gau2grid.git )
 set( GAUXC_GAU2GRID_REVISION   v2.0.6 )

include/gauxc/load_balancer.hpp

@@ -102,6 +102,12 @@ class LoadBalancer {
   const shell_pair_type& shell_pairs() const;
   const shell_pair_type& shell_pairs();
 
+  /// Return the underlying protonic BasisSet instance used to generate this LoadBalancer 
+  const basis_type& protonic_basis()  const;
+
+  /// Return the protonic BasisSetMap instance corresponding to protonic basis/molecule
+  const basis_map_type& protonic_basis_map() const;
+
   /// Return the runtime handle used to construct this LoadBalancer
   const RuntimeEnvironment& runtime() const;
 
@@ -163,6 +169,10 @@ class LoadBalancerFactory {
   LoadBalancer get_instance( const RuntimeEnvironment& rt, 
     const Molecule& mol, const MolGrid& mg, const BasisSet<double>& bs);
 
+  LoadBalancer get_instance( const RuntimeEnvironment& rt, 
+    const Molecule& mol, const MolGrid& mg, const BasisSet<double>& bs,
+    const BasisSet<double>& bs2 );
+
   /** 
    *  @brief Generate a shared pointer to a LoadBalancer instance per kernel and 
    *         execution space specfication
@@ -180,6 +190,11 @@ class LoadBalancerFactory {
     const RuntimeEnvironment& rt,
     const Molecule& mol, const MolGrid& mg, const BasisSet<double>&);
 
+  std::shared_ptr<LoadBalancer> get_shared_instance( 
+    const RuntimeEnvironment& rt,
+    const Molecule& mol, const MolGrid& mg, const BasisSet<double>&,
+    const BasisSet<double>& );
+
 private:
 
   ExecutionSpace ex_; ///< Execution space for the generated LoadBalancer instances

include/gauxc/xc_integrator.hpp

@@ -34,6 +34,8 @@ class XCIntegrator {
   using exc_vxc_type_rks  = std::tuple< value_type, matrix_type >;
   using exc_vxc_type_uks  = std::tuple< value_type, matrix_type, matrix_type >;  
   using exc_vxc_type_gks  = std::tuple< value_type, matrix_type, matrix_type, matrix_type, matrix_type >;
+  using exc_vxc_type_neo_rks = std::tuple< value_type, value_type, matrix_type, matrix_type, matrix_type >;  
+  using exc_vxc_type_neo_uks = std::tuple< value_type, value_type, matrix_type, matrix_type, matrix_type, matrix_type >;  
   using exc_grad_type = std::vector< value_type >;
   using exx_type      = matrix_type;
 
@@ -64,7 +66,11 @@ class XCIntegrator {
   exc_vxc_type_uks  eval_exc_vxc ( const MatrixType&, const MatrixType&,
                                    const IntegratorSettingsXC& = IntegratorSettingsXC{} );
   exc_vxc_type_gks  eval_exc_vxc ( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&,
-                                   const IntegratorSettingsXC& = IntegratorSettingsXC{});
+                                   const IntegratorSettingsXC& = IntegratorSettingsXC{} );
+  exc_vxc_type_neo_rks neo_eval_exc_vxc ( const MatrixType&, const MatrixType&, const MatrixType&, 
+                                          const IntegratorSettingsXC& = IntegratorSettingsXC{} );
+  exc_vxc_type_neo_uks neo_eval_exc_vxc ( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&,
+                                          const IntegratorSettingsXC& = IntegratorSettingsXC{} );
 
   exc_grad_type eval_exc_grad( const MatrixType& );
 

include/gauxc/xc_integrator/impl.hpp

@@ -70,9 +70,26 @@ template <typename MatrixType>
 typename XCIntegrator<MatrixType>::exc_vxc_type_gks
   XCIntegrator<MatrixType>::eval_exc_vxc( const MatrixType& Ps, const MatrixType& Pz, const MatrixType& Py, const MatrixType& Px, 
                                           const IntegratorSettingsXC& ks_settings ) {
-      if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
-        return pimpl_->eval_exc_vxc(Ps, Pz, Py, Px, ks_settings);
-  };
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->eval_exc_vxc(Ps, Pz, Py, Px, ks_settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_vxc_type_neo_rks
+  XCIntegrator<MatrixType>::neo_eval_exc_vxc( const MatrixType& elec_Ps, const MatrixType& prot_Ps, const MatrixType& prot_Pz,
+                                              const IntegratorSettingsXC& ks_settings ){
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->neo_eval_exc_vxc(elec_Ps, prot_Ps, prot_Pz, ks_settings);
+};
+
+template <typename MatrixType>
+typename XCIntegrator<MatrixType>::exc_vxc_type_neo_uks
+  XCIntegrator<MatrixType>::neo_eval_exc_vxc( const MatrixType& elec_Ps, const MatrixType& elec_Pz, const MatrixType& prot_Ps, const MatrixType& prot_Pz,
+                                              const IntegratorSettingsXC& ks_settings ){
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  return pimpl_->neo_eval_exc_vxc(elec_Ps, elec_Pz, prot_Ps, prot_Pz, ks_settings);
+};
+
 
 template <typename MatrixType>
 typename XCIntegrator<MatrixType>::exc_grad_type

include/gauxc/xc_integrator/integrator_factory.hpp

@@ -89,6 +89,58 @@ class XCIntegratorFactory {
     return get_shared_instance( std::make_shared<functional_type>(func), lb );
   }
 
+  /** Generate XCIntegrator instance
+   *
+   *  @param[in] func  XC functional
+   *  @param[in] epcfunc  EPC functional
+   *  @param[in] lb    Preconstructed Load Balancer instance
+   */
+  std::shared_ptr<integrator_type> get_shared_instance( 
+    std::shared_ptr<functional_type> func,
+    std::shared_ptr<functional_type> epcfunc,
+    std::shared_ptr<LoadBalancer>    lb ) {
+
+    // Create Local Work Driver
+    auto lwd = LocalWorkDriverFactory::make_local_work_driver( ex_, 
+      lwd_kernel_, local_work_settings_ );
+
+    // Create Reduction Driver
+    auto rd = ReductionDriverFactory::get_shared_instance( 
+      lb->runtime(), rd_kernel_ );
+
+    // Create Integrator instance
+    std::transform( input_type_.begin(), input_type_.end(), input_type_.begin(), 
+      ::toupper );
+
+    if(!epcfunc->is_polarized())
+      GAUXC_GENERIC_EXCEPTION("EPC FUNCTIONAL NOT POLARIZED");
+
+    if( input_type_ == "REPLICATED" )
+      return std::make_shared<integrator_type>( 
+        ReplicatedXCIntegratorFactory<MatrixType>::make_integrator_impl(
+          ex_, integrator_kernel_, func, epcfunc, lb, std::move(lwd), rd
+        )
+      );
+    else GAUXC_GENERIC_EXCEPTION("INTEGRATOR TYPE NOT RECOGNIZED");
+
+    return nullptr;
+
+  }
+
+  auto get_shared_instance( const functional_type& func, const functional_type& epcfunc,
+                            const LoadBalancer& lb ) {
+    return get_shared_instance( std::make_shared<functional_type>(func),
+                                std::make_shared<functional_type>(epcfunc),
+                                std::make_shared<LoadBalancer>(lb) );
+  }
+
+  auto get_shared_instance( const functional_type& func, const functional_type& epcfunc,
+                            std::shared_ptr<LoadBalancer> lb ) {
+    return get_shared_instance( std::make_shared<functional_type>(func), 
+                                std::make_shared<functional_type>(epcfunc),
+                                lb );
+  }
+
 
   template <typename... Args>
   integrator_type get_instance( Args&&... args ) {

include/gauxc/xc_integrator/replicated/impl.hpp

@@ -157,6 +157,59 @@ typename ReplicatedXCIntegrator<MatrixType>::exc_vxc_type_gks
 
 }
 
+template <typename MatrixType>
+typename ReplicatedXCIntegrator<MatrixType>::exc_vxc_type_neo_rks
+  ReplicatedXCIntegrator<MatrixType>::neo_eval_exc_vxc_( const MatrixType& elec_Ps, const MatrixType& prot_Ps, const MatrixType& prot_Pz,
+                                                         const IntegratorSettingsXC& ks_settings ) {
+
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  matrix_type elec_VXCs( elec_Ps.rows(), elec_Ps.cols() );
+  matrix_type prot_VXCs( prot_Ps.rows(), prot_Ps.cols() );
+  matrix_type prot_VXCz( prot_Pz.rows(), prot_Pz.cols() );
+  value_type  elec_EXC;
+  value_type  prot_EXC;
+
+  pimpl_->neo_eval_exc_vxc( elec_Ps.rows(), elec_Ps.cols(), prot_Ps.rows(), prot_Ps.cols(),
+                            elec_Ps.data(), elec_Ps.rows(),
+                            prot_Ps.data(), prot_Ps.rows(),
+                            prot_Pz.data(), prot_Pz.rows(),
+                            elec_VXCs.data(), elec_VXCs.rows(),
+                            prot_VXCs.data(), prot_VXCs.rows(),
+                            prot_VXCz.data(), prot_VXCz.rows(),
+                            &elec_EXC, &prot_EXC, ks_settings );
+
+  return std::make_tuple( elec_EXC, prot_EXC, elec_VXCs, prot_VXCs, prot_VXCz );
+
+}
+
+template <typename MatrixType>
+typename ReplicatedXCIntegrator<MatrixType>::exc_vxc_type_neo_uks
+  ReplicatedXCIntegrator<MatrixType>::neo_eval_exc_vxc_( const MatrixType& elec_Ps, const MatrixType& elec_Pz, const MatrixType& prot_Ps, const MatrixType& prot_Pz,
+                                                         const IntegratorSettingsXC& ks_settings ) {
+
+  if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
+  matrix_type elec_VXCs( elec_Ps.rows(), elec_Ps.cols() );
+  matrix_type elec_VXCz( elec_Pz.rows(), elec_Pz.cols() );
+  matrix_type prot_VXCs( prot_Ps.rows(), prot_Ps.cols() );
+  matrix_type prot_VXCz( prot_Pz.rows(), prot_Pz.cols() );
+  value_type  elec_EXC;
+  value_type  prot_EXC;
+
+  pimpl_->neo_eval_exc_vxc( elec_Ps.rows(), elec_Ps.cols(), prot_Ps.rows(), prot_Ps.cols(),
+                            elec_Ps.data(), elec_Ps.rows(),
+                            elec_Pz.data(), elec_Pz.rows(),
+                            prot_Ps.data(), prot_Ps.rows(),
+                            prot_Pz.data(), prot_Pz.rows(),
+                            elec_VXCs.data(), elec_VXCs.rows(),
+                            elec_VXCz.data(), elec_VXCz.rows(),
+                            prot_VXCs.data(), prot_VXCs.rows(),
+                            prot_VXCz.data(), prot_VXCz.rows(),
+                            &elec_EXC, &prot_EXC, ks_settings );
+
+  return std::make_tuple( elec_EXC, prot_EXC, elec_VXCs, elec_VXCz, prot_VXCs, prot_VXCz );
+
+}
+
 template <typename MatrixType>
 typename ReplicatedXCIntegrator<MatrixType>::exc_grad_type 
   ReplicatedXCIntegrator<MatrixType>::eval_exc_grad_( const MatrixType& P ) {

include/gauxc/xc_integrator/replicated/replicated_xc_device_integrator.hpp

@@ -57,6 +57,15 @@ struct ReplicatedXCDeviceIntegratorFactory {
     std::shared_ptr<ReductionDriver>   rd
     );
 
+  static ptr_return_t make_integrator_impl( 
+    std::string integrator_kernel,
+    std::shared_ptr<functional_type>   func,
+    std::shared_ptr<functional_type>   epcfunc,
+    std::shared_ptr<LoadBalancer>      lb,
+    std::unique_ptr<LocalWorkDriver>&& lwd, 
+    std::shared_ptr<ReductionDriver>   rd
+    );
+
 };
 
 

include/gauxc/xc_integrator/replicated/replicated_xc_host_integrator.hpp

@@ -56,6 +56,15 @@ struct ReplicatedXCHostIntegratorFactory {
     std::unique_ptr<LocalWorkDriver>&& lwd,
     std::shared_ptr<ReductionDriver>   rd
     );
+
+  static ptr_return_t make_integrator_impl( 
+    std::string integrator_kernel,
+    std::shared_ptr<functional_type>   func,
+    std::shared_ptr<functional_type>   epcfunc,
+    std::shared_ptr<LoadBalancer>      lb,
+    std::unique_ptr<LocalWorkDriver>&& lwd,
+    std::shared_ptr<ReductionDriver>   rd
+    );
 
 };
 

include/gauxc/xc_integrator/replicated/replicated_xc_integrator_factory.hpp

@@ -73,6 +73,59 @@ struct ReplicatedXCIntegratorFactory {
 
   }
 
+
+
+  /** Generate a ReplicatedXCIntegrator instance
+   *
+   *  @param[in]  ex                 Execution space for integrator instance
+   *  @param[in]  integration_kernel Name of integration scaffold to load ("Default", "Reference", etc)
+   *  @param[in]  func               XC functional to integrate
+   *  @param[in]  epcfunc            EPC functional to integrate
+   *  @param[in]  lb                 Pregenerated LoadBalancer instance
+   *  @param[in]  lwd                Local Work Driver
+   */
+  static ptr_return_t make_integrator_impl( 
+    ExecutionSpace ex,
+    std::string integrator_kernel,
+    std::shared_ptr<functional_type>   func,
+    std::shared_ptr<functional_type>   epcfunc,
+    std::shared_ptr<LoadBalancer>      lb,
+    std::unique_ptr<LocalWorkDriver>&& lwd,
+    std::shared_ptr<ReductionDriver>   rd
+    ) {
+
+
+
+    switch(ex) {
+
+      using host_factory = 
+        detail::ReplicatedXCHostIntegratorFactory<value_type>;
+      case ExecutionSpace::Host:
+        return std::make_unique<integrator_type>( 
+          host_factory::make_integrator_impl(
+            integrator_kernel, func, epcfunc, lb, std::move(lwd), rd 
+          )
+        );
+
+      #ifdef GAUXC_HAS_DEVICE
+      using device_factory = 
+        detail::ReplicatedXCDeviceIntegratorFactory<value_type>;
+      case ExecutionSpace::Device:
+        return std::make_unique<integrator_type>( 
+          device_factory::make_integrator_impl(
+            integrator_kernel, func, epcfunc, lb, std::move(lwd), rd
+          )
+        );
+      #endif
+
+      default:
+        GAUXC_GENERIC_EXCEPTION("ReplicatedXCIntegrator ExecutionSpace Not Supported");
+    }
+
+    return nullptr;
+
+  }
+
 
 };
 

include/gauxc/xc_integrator/replicated/replicated_xc_integrator_impl.hpp

@@ -29,6 +29,7 @@ class ReplicatedXCIntegratorImpl {
 protected:
 
   std::shared_ptr< functional_type > func_;               ///< XC functional
+  std::shared_ptr< functional_type > epcfunc_;            ///< EPC functional for NEO
   std::shared_ptr< LoadBalancer >    load_balancer_;      ///< Load Balancer
   std::unique_ptr< LocalWorkDriver > local_work_driver_;  ///< Local Work Driver
   std::shared_ptr< ReductionDriver > reduction_driver_;   ///< Reduction Driver
@@ -73,7 +74,24 @@ class ReplicatedXCIntegratorImpl {
                               value_type* VXCy, int64_t ldvxcy,
                               value_type* VXCx, int64_t ldvxcx,
                               value_type* EXC, const IntegratorSettingsXC& ks_settings ) = 0;
-
+  virtual void neo_eval_exc_vxc_( int64_t elec_m, int64_t elec_mn, int64_t prot_m, int64_t prot_n,
+                                  const value_type* elec_Ps, int64_t elec_ldps,
+                                  const value_type* prot_Ps, int64_t prot_ldps,
+                                  const value_type* prot_Pz, int64_t prot_ldpz,
+                                  value_type* elec_VXCs,     int64_t elec_ldvxcs,
+                                  value_type* prot_VXCs,     int64_t prot_ldvxcs,
+                                  value_type* prot_VXCz,     int64_t prot_ldvxcz,
+                                  value_type* elec_EXC,  value_type* prot_EXC, const IntegratorSettingsXC& ks_settings) = 0;
+  virtual void neo_eval_exc_vxc_( int64_t elec_m, int64_t elec_n, int64_t prot_m, int64_t prot_n, 
+                                  const value_type* elec_Ps, int64_t elec_ldps,
+                                  const value_type* elec_Pz, int64_t elec_ldpz,
+                                  const value_type* prot_Ps, int64_t prot_ldps,
+                                  const value_type* prot_Pz, int64_t prot_ldpz,
+                                  value_type* elec_VXCs,     int64_t elec_ldvxcs,
+                                  value_type* elec_VXCz,     int64_t elec_ldvxcz,
+                                  value_type* prot_VXCs,     int64_t prot_ldvxcs,
+                                  value_type* prot_VXCz,     int64_t prot_ldvxcz,
+                                  value_type* elec_EXC,  value_type* prot_EXC, const IntegratorSettingsXC& ks_settings) = 0;
   virtual void eval_exc_grad_( int64_t m, int64_t n, const value_type* P,
                                int64_t ldp, value_type* EXC_GRAD ) = 0;
   virtual void eval_exx_( int64_t m, int64_t n, const value_type* P,
@@ -88,6 +106,13 @@ class ReplicatedXCIntegratorImpl {
                               std::shared_ptr< ReductionDriver>    rd
                               );
 
+  ReplicatedXCIntegratorImpl( std::shared_ptr< functional_type >   func,
+                              std::shared_ptr< functional_type >   epcfunc,
+                              std::shared_ptr< LoadBalancer >      lb, 
+                              std::unique_ptr< LocalWorkDriver >&& lwd,
+                              std::shared_ptr< ReductionDriver>    rd
+                              );
+
   virtual ~ReplicatedXCIntegratorImpl() noexcept;
 
   void integrate_den( int64_t m, int64_t n, const value_type* P,
@@ -128,7 +153,26 @@ class ReplicatedXCIntegratorImpl {
                      value_type* VXCy, int64_t ldvxcy,
                      value_type* VXCx, int64_t ldvxcx,
                      value_type* EXC, const IntegratorSettingsXC& ks_settings );
-
+
+  void neo_eval_exc_vxc( int64_t elec_m, int64_t elec_n, int64_t prot_m, int64_t prot_n,  
+                         const value_type* elec_Ps, int64_t elec_ldps,
+                         const value_type* prot_Ps, int64_t prot_ldps,
+                         const value_type* prot_Pz, int64_t prot_ldpz,
+                         value_type* elec_VXCs,     int64_t elec_ldvxcs,
+                         value_type* prot_VXCs,     int64_t prot_ldvxcs,
+                         value_type* prot_VXCz,     int64_t prot_ldvxcz,
+                         value_type* elec_EXC,  value_type* prot_EXC, const IntegratorSettingsXC& ks_settings );
+
+  void neo_eval_exc_vxc( int64_t elec_m, int64_t elec_n, int64_t prot_m, int64_t prot_n, 
+                         const value_type* elec_Ps, int64_t elec_ldps,
+                         const value_type* elec_Pz, int64_t elec_ldpz,
+                         const value_type* prot_Ps, int64_t prot_ldps,
+                         const value_type* prot_Pz, int64_t prot_ldpz,
+                         value_type* elec_VXCs,     int64_t elec_ldvxcs,
+                         value_type* elec_VXCz,     int64_t elec_ldvxcz,
+                         value_type* prot_VXCs,     int64_t prot_ldvxcs,
+                         value_type* prot_VXCz,     int64_t prot_ldvxcz,
+                         value_type* elec_EXC,  value_type* prot_EXC, const IntegratorSettingsXC& ks_settings );
 
   void eval_exc_grad( int64_t m, int64_t n, const value_type* P,
                       int64_t ldp, value_type* EXC_GRAD );