Merge pull request #118 from adc-connect/libxm

Adapt AdcMemory for libxm

.github/workflows/ci.yaml

@@ -48,15 +48,24 @@ jobs:
         run: |
           mkdir $HOME/.adcc/
           echo 'coverage = True' > $HOME/.adcc/siteconfig.py
+      - name: Hard-code download url for MacOS
+        run: |
+          echo "libtensor_url = \"$BASE_URL/$VERSION_URL\"" >> $HOME/.adcc/siteconfig.py
+        env:
+          BASE_URL: https://github.com/adc-connect/libtensor/releases/download
+          VERSION_URL: v3.0.1/libtensorlight-3.0.1-macosx_10_15_x86_64.tar.gz
+        if: contains( matrix.os, 'macos')
       - name: Install python dependencies
         run: |
           export PATH="/usr/local/opt/ccache/libexec:$PATH"
           export PATH="/usr/lib/ccache:$PATH"
           python -m pip install --user pyscf cppe wheel
           python -m pip install --user -r requirements.txt -v
       #
-      - name: Run python tests
+      - name: Run python tests with std allocator
         run: python setup.py test -a '--cov=adcc'
+      - name: Run reduced python tests with libxm
+        run: python setup.py test -a '--allocator=libxm -k "TestFunctionality and h2o_sto3g"'
       - name: Run C++ tests
         run: python setup.py cpptest -v
       #

adcc/conftest.py

@@ -21,7 +21,6 @@
 ##
 ## ---------------------------------------------------------------------
 import os
-
 import pytest
 
 from .testdata.cache import TestdataCache
@@ -61,6 +60,10 @@ def pytest_addoption(parser):
         "--skip-update", default=False, action="store_true",
         help="Skip updating testdata"
     )
+    parser.addoption(
+        "--allocator", default="standard", choices=["standard", "libxm"],
+        help="Allocator to use for the tests"
+    )
 
 
 def pytest_collection_modifyitems(config, items):
@@ -85,3 +88,9 @@ def pytest_collection(session):
 def pytest_runtestloop(session):
     if os.environ.get("CI", "false") == "true":
         setup_continuous_integration(session)
+    if session.config.option.allocator != "standard":
+        import adcc
+
+        allocator = session.config.option.allocator
+        adcc.memory_pool.initialise(allocator=allocator)
+        print(f"Using allocator: {allocator}")

adcc/memory_pool.py

@@ -29,44 +29,42 @@
 
 
 class MemoryPool(libadcc.AdcMemory):
-    def initialise(self, max_memory, tensor_block_size=16,
-                   pagefile_directory=None, allocator="default"):
+    def initialise(self, scratch_directory="/tmp", max_block_size=16,
+                   allocator="standard"):
         """Initialise the adcc memory management.
 
         Parameters
         ----------
-        max_memory : int
-            Estimate for the maximally employed memory. Note that this value
-            is only effective if the allocator parameter is "libxm". and not
-            for "default" or "standard".
+        scratch_directory : str, optional
+            Directory for storing temporary pagefiles. This should be a fast
+            storage location as tensor data will be mapped to asynchronously
+            to this directory for the "libxc" allocator.
 
-        tensor_block_size : int, optional
-            This parameter roughly has the meaning of how many indices are handled
-            together on operations. A good value is 16 for most nowaday CPU
-            cachelines.
-
-        pagefile_prefix : str, optional
-            Directory prefix for storing temporary cache files.
+        max_block_size : int, optional
+            The maximal size a tensor block may have along each axis.
 
         allocator : str, optional
-            The allocator to be used. Valid values are "libxm", "standard"
-            (libstc++ allocator) and "default", where "default" uses the
-            best-available default.
+            The allocator to be used. Valid values are "libxm" or "standard"
+            (libstc++ allocator).
         """
-        if not pagefile_directory:
-            pagefile_directory = tempfile.mkdtemp(prefix="adcc_", dir="/tmp")
-        super().initialise(pagefile_directory, max_memory, tensor_block_size,
-                           allocator)
+        pagefile_directory = tempfile.mkdtemp(prefix="adcc_", dir=scratch_directory)
+        super().initialise(pagefile_directory, max_block_size, allocator)
         atexit.register(MemoryPool.cleanup, self)
 
     def cleanup(self):
         if os.path.isdir(self.pagefile_directory):
             shutil.rmtree(self.pagefile_directory)
 
+    @property
+    def scratch_directory(self):
+        return os.path.dirname(self.pagefile_directory)
+
     @property
     def page_files(self):
         """The list of all page files."""
-        return glob.glob(os.path.join(self.pagefile_directory, "pagefile.*"))
+        globs = ["pagefile.*", "xmpagefile"]
+        globs = [os.path.join(self.pagefile_directory, pat) for pat in globs]
+        return [c for pat in globs for c in glob.glob(pat)]
 
     @property
     def total_size_page_files(self):

conda/meta.yaml.in

@@ -10,10 +10,10 @@ requirements:
     - {{ compiler('c') }}
     - {{ compiler('cxx') }}
     - python {{ python }}
-    - libtensorlight
+    - libtensorlight >=3.0.1
   host:
     - python {{ python }}
-    - libtensorlight
+    - libtensorlight >=3.0.1
     - pybind11 >=2.6
     # Testing:
     - pytest
@@ -31,7 +31,7 @@ requirements:
     - pandas >=0.25.0
   run:
     - {{ pin_compatible('python', max_pin='x.x') }}
-    - libtensorlight
+    - libtensorlight >=3.0.1
     - numpy >=1.14
     - h5py >=2.9
     - scipy >=1.2

libadcc/AdcMemory.cc

@@ -20,103 +20,62 @@
 #include "AdcMemory.hh"
 #include "exceptions.hh"
 #include <algorithm>
-#include <libtensor/metadata.h>
 
 // Change visibility of libtensor singletons to public
 #pragma GCC visibility push(default)
 #include <libtensor/core/allocator.h>
 #include <libtensor/core/batching_policy_base.h>
+#include <libtensor/expr/btensor/eval_btensor.h>
+#include <libtensor/metadata.h>
 #pragma GCC visibility pop
 
 namespace libadcc {
+namespace {
+bool libtensor_has_libxm() {
+  const std::vector<std::string>& libtensor_features = libtensor::metadata::features();
+  const auto it =
+        std::find(libtensor_features.begin(), libtensor_features.end(), "libxm");
+  return it != libtensor_features.end();
+}
+}  // namespace
 
 AdcMemory::AdcMemory()
-      : m_allocator("none"),
+      : m_max_block_size(0),
+        m_allocator("none"),
         m_initialise_called(false),
-        m_max_memory(0),
-        m_pagefile_directory{""},
-        m_tbs_param(0) {
-  // Just some value: >4GB is quite usual today and should get the setup right.
-  // ... and anyway this limit is not really a hard limit for the std::allocator
-  // anyway, since it plainly allocates memory on the system until we run out.
-  const size_t max_memory = 4ul * 1024ul * 1024ul * 1024ul;
-
-  // A tbs_param of 16 should be a good default.
-  const size_t tbs_param = 16;
-
-  initialise("", max_memory, tbs_param, "standard");
+        m_pagefile_directory{""} {
+  initialise("", 16, "standard");
 
   // Make sure we can call initialise once more
   m_initialise_called = false;
 }
 
-void AdcMemory::initialise(std::string pagefile_directory, size_t max_memory,
-                           size_t tbs_param, std::string allocator) {
-  if (max_memory == 0) {
-    throw invalid_argument("A max_memory value of 0 is not valid.");
-  }
-  if (tbs_param == 0) {
-    throw invalid_argument("A tbs_param value of 0 is not valid.");
-  }
+void AdcMemory::initialise(std::string pagefile_directory, size_t max_block_size,
+                           std::string allocator) {
   if (m_initialise_called) {
     throw invalid_argument("Cannot initialise AdcMemory object twice.");
   }
 
   m_initialise_called  = true;
-  m_max_memory         = max_memory;
   m_pagefile_directory = pagefile_directory;
-  m_tbs_param          = tbs_param;
-
-  //
-  // Determine initialisation parameters
-  //
-  // Our principle unit of memory is the employed scalar type,
-  // so our memory calculations will be based on that.
-  const size_t memunit = sizeof(scalar_type);
-
-  // Ideally our largest tensor block should be large enough to fit a few of the
-  // block size parameters. Our rationale here is that we have typically
-  // no more than tensors of rank 6
-  const size_t tbs3               = tbs_param * tbs_param * tbs_param;
-  const size_t largest_block_size = tbs3 * tbs3;
-
-  if (largest_block_size * memunit > max_memory) {
-    throw invalid_argument("At least " + std::to_string(largest_block_size * memunit) +
-                           " bytes of memory need to be requested.");
-  }
-
-  const bool has_libxm = [] {
-    const std::vector<std::string>& libtensor_features = libtensor::metadata::features();
-    const auto it =
-          std::find(libtensor_features.begin(), libtensor_features.end(), "libxm");
-    return it != libtensor_features.end();
-  }();
-
-  if (allocator == "default") {
-    // Preference is 1. libxm 2. standard
-    allocator = has_libxm ? "libxm" : "standard";
-  }
+  m_max_block_size     = max_block_size;
 
   if (allocator == "standard") {
-  } else if (allocator == "libxm" && has_libxm) {
+  } else if (allocator == "libxm" && libtensor_has_libxm()) {
   } else {
     throw invalid_argument("A memory allocator named '" + allocator +
                            "' is not known to adcc. Perhaps it is not compiled in.");
   }
 
   shutdown();  // Shutdown previously initialised allocator (if any)
-  libtensor::allocator<double>::init(
-        allocator,
-        tbs_param,                     // Exponential base of data block size
-        tbs_param * memunit,           // Smallest data block size in bytes
-        largest_block_size * memunit,  // Largest data block size in bytes
-        max_memory,                    // Memory limit in bytes
-        pagefile_directory.c_str()     // Prefix to page file path.
-  );
+  libtensor::allocator<double>::init(allocator, pagefile_directory.c_str());
   m_allocator = allocator;
 
-  // Set initial batch size
-  set_contraction_batch_size(max_memory / tbs3 / tbs_param / 3);
+  // Set initial batch size ... value empirically determined to be reasonable
+  set_contraction_batch_size(21870);
+
+  // Enable libxm contraction backend if libxm allocator is used
+  libtensor::expr::eval_btensor<double>::use_libxm(allocator == "libxm");
 }
 
 size_t AdcMemory::contraction_batch_size() const {

libadcc/AdcMemory.hh

@@ -39,10 +39,9 @@ namespace libadcc {
  */
 class AdcMemory {
  public:
-  /** Setup the class and directly call initialise with the default allocator */
-  AdcMemory(std::string pagefile_directory, size_t max_memory, size_t tbs_param = 16)
-        : AdcMemory{} {
-    initialise(pagefile_directory, max_memory, tbs_param);
+  /** Setup the class and directly call initialise with the standard allocator */
+  AdcMemory(std::string pagefile_directory, size_t max_block_size = 16) : AdcMemory{} {
+    initialise(pagefile_directory, max_block_size);
   }
 
   /** Setup the AdcMemory class using the std::allocator for memory management */
@@ -58,65 +57,50 @@ class AdcMemory {
   /** Return the allocator to which the class is initialised. */
   std::string allocator() const { return m_allocator; }
 
-  /** Return the max_memory parameter value to which the class was initialised.
-   *  \note This value is only a meaningful upper bound if
-   *  allocator() != "standard" */
-  size_t max_memory() const { return m_max_memory; }
-
   /** Return the pagefileprefix value
    *
    * \note This value is only meaningful if allocator() != "standard" */
   std::string pagefile_directory() const { return m_pagefile_directory; }
 
-  /** Return the tbs_param value */
-  size_t tbs_param() const { return m_tbs_param; }
-
   /** Get the contraction batch size, this is the number of
-   *  tensor blocks, which are processed in a batch in a tensor contraction. */
+   *  tensor elements, which are processed in a batch in a tensor contraction. */
   size_t contraction_batch_size() const;
 
   /** Set the contraction batch size */
   void set_contraction_batch_size(size_t bsize);
 
+  /** Get the maximal block size a tensor may have along each axis */
+  size_t max_block_size() const { return m_max_block_size; }
+
   /** Setup the environment for the memory management.
    *
    * \param pagefile_directory  File prefix for page files
-   * \param max_memory          The maximal memory adc makes use of (in bytes).
-   * \param tbs_param           The tensor block size parameter.
-   *                            This parameter roughly has the meaning of how many indices
-   *                            are handled together on operations. A good value seems to
-   *                            be 16.
-   * \param allocator   The allocator to be used. Valid values are "libxm", "libvmm",
-   *                    "standard" and "default", where "default" uses a default
-   *                    chosen from the first three.
+   * \param max_block_size      Maximal block size a tensor may have along each axis.
+   * \param allocator   The allocator to be used. Valid values are "libxm" or
+   *                    "standard" where "standard" is just the standard C++ allocator.
    *
-   * \note For some allocators \c pagefile_directory and \c max_memory are not supported
-   *       and thus ignored.
-   * \note "libxm" and "libvmm" are extra features, which are not available in a
-   *       default setup.
+   * \note For some allocators \c pagefile_directory is not supported and thus ignored.
+   * \note "libxm" is an extra features, which might not be available in a default setup.
    **/
-  void initialise(std::string pagefile_directory, size_t max_memory,
-                  size_t tbs_param = 16, std::string allocator = "default");
+  void initialise(std::string pagefile_directory, size_t max_block_size = 16,
+                  std::string allocator = "standard");
 
  protected:
   /** Shutdown the allocator, i.e. cleanup all memory currently held. */
   void shutdown();
 
  private:
+  /** Maximal size a tensor block may have along any axis. */
+  size_t m_max_block_size;
+
   /** The allocator this object is currently initialised to. */
   std::string m_allocator;
 
   /** Has the initialise function been called by the user */
   bool m_initialise_called;
 
-  /** Configured maximal memory in bytes */
-  size_t m_max_memory;
-
   /** Configured file prefix for pagefiles */
   std::string m_pagefile_directory;
-
-  /** Configured tensor block size parameter */
-  size_t m_tbs_param;
 };
 
 ///@}

libadcc/MoSpaces.cc

@@ -446,7 +446,7 @@ MoSpaces::MoSpaces(const HartreeFockSolution_i& hf,
   //
   // Lambda to construct the blocks for a particular space.
   // Modifies map_block_start, map_block_spin, map_block_irrep
-  const size_t max_block_size = adcmem_ptr->tbs_param();
+  const size_t max_block_size = adcmem_ptr->max_block_size();
   auto construct_blocks_for   = [this, &map_index_sub_to_full, &irrep_of, &spin_of,
                                &subspace_of, &max_block_size](const std::string& space) {
     std::vector<size_t> fidx_of;

libadcc/Symmetry.cc

@@ -118,10 +118,9 @@ Symmetry::Symmetry(std::shared_ptr<const MoSpaces> mospaces_ptr, const std::stri
                                 mospaces_ptr->map_block_start.at(ss)});
     } else {
       // ss is an extra axes
-      const size_t max_block_size = mospaces_ptr->adcmem_ptr()->tbs_param();
-      const size_t n_orbs_alpha   = itextra->second.first;
-      const size_t n_orbs_beta    = itextra->second.second;
-      const size_t n_orbs         = n_orbs_alpha + n_orbs_beta;
+      const size_t n_orbs_alpha = itextra->second.first;
+      const size_t n_orbs_beta  = itextra->second.second;
+      const size_t n_orbs       = n_orbs_alpha + n_orbs_beta;
 
       if (n_orbs_alpha == 0) {
         throw invalid_argument(
@@ -135,8 +134,8 @@ Symmetry::Symmetry(std::shared_ptr<const MoSpaces> mospaces_ptr, const std::stri
         splits_crude.push_back(n_orbs_alpha);
       }
 
-      const std::vector<size_t> blks =
-            construct_blocks(splits_crude, n_orbs, max_block_size);
+      const std::vector<size_t> blks = construct_blocks(
+            splits_crude, n_orbs, mospaces_ptr->adcmem_ptr()->max_block_size());
 
       m_axes.push_back(AxisInfo{ss, n_orbs_alpha, n_orbs_beta, blks});
     }