refactor(lcao): decouple basis indexing from block-cyclic layouts

examples/13_dos/01_pw_Al/run.sh

@@ -14,7 +14,7 @@ OMP_NUM_THREADS=${ABACUS_THREADS} mpirun -np ${ABACUS_NPROCS} ${ABACUS_PATH} | t
 
 rm INPUT KPT
 
-if [[ ! -f scf.output ]] || 
+if [[ ! -f scf.output ]] ||
    [[ ! -f nscf.output ]] ||
    [[ ! -f OUT.ABACUS/running_scf.log ]] ||
    [[ ! -f OUT.ABACUS/running_nscf.log ]] ||
@@ -28,4 +28,4 @@ then
 else
 	echo "job succeeded!"
 	exit 0
-fi
+fi

examples/13_dos/02_lcao_Si2/run.sh

@@ -14,7 +14,7 @@ OMP_NUM_THREADS=${ABACUS_THREADS} mpirun -np ${ABACUS_NPROCS} ${ABACUS_PATH} | t
 
 rm INPUT KPT
 
-if [[ ! -f scf.output ]] || 
+if [[ ! -f scf.output ]] ||
    [[ ! -f nscf.output ]] ||
    [[ ! -f OUT.ABACUS/running_scf.log ]] ||
    [[ ! -f OUT.ABACUS/running_nscf.log ]] ||
@@ -28,4 +28,4 @@ then
 else
 	echo "job succeeded!"
 	exit 0
-fi
+fi

source/Makefile.Objects

@@ -153,6 +153,8 @@ OBJS_BASE=assoc_laguerre.o\
     complexarray.o\
     complexmatrix.o\
     clebsch_gordan_coeff.o\
+    basis_flattened_indexing.o\
+    atom_basis_view.o\
     element_basis_index.o\
     inverse_matrix.o\
     gather_math_lib_info.o\
@@ -200,7 +202,7 @@ OBJS_BASE=assoc_laguerre.o\
     device.o\
     device_helpers.o\
     output_device.o\
-    block_cyclic_distribution.o\
+    parallel_blkcyc.o\
     binstream.o\
     filename.o\
     sparse_matrix.o\

source/source_base/CMakeLists.txt

@@ -3,83 +3,87 @@ if(ENABLE_ABACUS_LIBM)
        libm/sincos.cpp)
 endif()
 add_library(
-  base OBJECT
-  assoc_laguerre.cpp
-  module_external/blas_connector_base.cpp
-  module_external/blas_connector_vector.cpp
-  module_external/blas_connector_matrix.cpp
-  clebsch_gordan_coeff.cpp
-  complexarray.cpp
-  complexmatrix.cpp
-  element_basis_index.cpp
-  inverse_matrix.cpp
-  gather_math_lib_info.cpp
-  global_file.cpp
-  global_function.cpp
-  glob_fn_ddotreal.cpp
-  global_variable.cpp
-  intarray.cpp
-  math_integral.cpp
-  math_lebedev_laikov.cpp
-  math_polyint.cpp
-  math_sphbes.cpp
-  math_erf_complex.cpp
-  math_ylmreal.cpp
-  math_bspline.cpp
-  math_chebyshev.cpp
-  mathzone_add1.cpp
-  matrix.cpp
-  matrix3.cpp
-  memory_recorder.cpp
-  mymath.cpp
-  opt_cg.cpp
-  opt_dcsrch.cpp
-  opt_tn.cpp
-  output.cpp
-  para_gemm.cpp
-  realarray.cpp
-  sph_bessel_rec_d1.cpp
-  sph_bessel_rec_d2.cpp
-  timer.cpp
-  tool_check.cpp
-  tool_quit.cpp
-  tool_title.cpp
-  ylm.cpp
-  parallel_common.cpp
-  parallel_cell.cpp
-  parallel_global.cpp
-  parallel_comm.cpp
-  parallel_reduce.cpp
-  parallel_device.cpp
-  parallel_grid.cpp
-  sph_bessel_tf.cpp
-  cubic_spline.cpp
-  block_cyclic_distribution.cpp
-  module_out/binstream.cpp
-  module_out/filename.cpp
-  module_out/sparse_matrix.cpp
-  module_out/file_reader.cpp
-  module_out/csr_reader.cpp
-  module_out/read_exit_file.cpp
-  module_out/orb_io.cpp
-  module_mixing/mixing_data.cpp
-  module_mixing/mixing.cpp
-  module_mixing/plain_mixing.cpp
-  module_mixing/pulay_mixing.cpp
-  module_mixing/broyden_mixing.cpp
-  module_parallel/para_world.cpp
-  module_parallel/para_band_output.cpp
-  module_parallel/para_collection.cpp
-  module_parallel/para_kmesh_world.cpp
-  module_parallel/para_pw_world.cpp
-  module_parallel/para_diag_world.cpp
-  module_parallel/para_rgrid_world.cpp
-  module_parallel/para_bgroup_world.cpp
-  module_parallel/para_matrix_world.cpp
-  module_parallel/para_mpi_func.cpp
-  module_parallel/para_setup.cpp
-  module_parallel/para_bridge.cpp
-  ${LIBM_SRC})
+    base
+    OBJECT
+    assoc_laguerre.cpp
+    atom_basis_view.cpp
+    module_external/blas_connector_base.cpp
+    module_external/blas_connector_vector.cpp
+    module_external/blas_connector_matrix.cpp
+    clebsch_gordan_coeff.cpp
+    complexarray.cpp
+    complexmatrix.cpp
+    basis_flattened_indexing.cpp
+    element_basis_index.cpp
+    inverse_matrix.cpp
+    gather_math_lib_info.cpp
+    global_file.cpp
+    global_function.cpp
+    glob_fn_ddotreal.cpp
+    global_variable.cpp
+    intarray.cpp
+    math_integral.cpp
+    math_lebedev_laikov.cpp
+    math_polyint.cpp
+    math_sphbes.cpp
+    math_erf_complex.cpp
+    math_ylmreal.cpp
+    math_bspline.cpp
+    math_chebyshev.cpp
+    mathzone_add1.cpp
+    matrix.cpp
+    matrix3.cpp
+    memory_recorder.cpp
+    mymath.cpp
+    opt_cg.cpp
+    opt_dcsrch.cpp
+    opt_tn.cpp
+    output.cpp
+    para_gemm.cpp
+    realarray.cpp
+    sph_bessel_rec_d1.cpp
+    sph_bessel_rec_d2.cpp
+    timer.cpp
+    tool_check.cpp
+    tool_quit.cpp
+    tool_title.cpp
+    ylm.cpp
+    parallel_common.cpp
+    parallel_cell.cpp
+    parallel_global.cpp
+    parallel_comm.cpp
+    parallel_reduce.cpp
+    parallel_device.cpp
+    parallel_grid.cpp
+    sph_bessel_tf.cpp
+    cubic_spline.cpp
+    parallel_blkcyc.cpp
+    module_out/binstream.cpp
+    module_out/filename.cpp
+    module_out/sparse_matrix.cpp
+    module_out/file_reader.cpp
+    module_out/csr_reader.cpp
+    module_out/read_exit_file.cpp
+    module_out/orb_io.cpp
+    module_mixing/mixing_data.cpp
+    module_mixing/mixing.cpp
+    module_mixing/plain_mixing.cpp
+    module_mixing/pulay_mixing.cpp
+    module_mixing/broyden_mixing.cpp
+    module_parallel/para_world.cpp
+    module_parallel/para_band_output.cpp
+    module_parallel/para_collection.cpp
+    module_parallel/para_kmesh_world.cpp
+    module_parallel/para_pw_world.cpp
+    module_parallel/para_diag_world.cpp
+    module_parallel/para_rgrid_world.cpp
+    module_parallel/para_bgroup_world.cpp
+    module_parallel/para_matrix_world.cpp
+    module_parallel/para_mpi_func.cpp
+    module_parallel/para_setup.cpp
+    module_parallel/para_bridge.cpp
+    ${LIBM_SRC}
+)
 
 target_link_libraries(base PUBLIC container)
 add_subdirectory(module_container)
@@ -89,19 +93,36 @@ if(ENABLE_COVERAGE)
 endif()
 
 if(BUILD_TESTING)
-  # TODO: classify the serial and parrallel tests
   AddTest(
-    TARGET MODULE_BASE_block_cyclic_distribution
+    TARGET MODULE_BASE_atom_basis_view
+    LIBS parameter
+    SOURCES test/test_atom_basis_view.cpp
+            atom_basis_view.cpp
+            basis_flattened_indexing.cpp
+            parallel_blkcyc.cpp
+            element_basis_index.cpp
+  )
+  AddTest(
+    TARGET MODULE_BASE_parallel_blkcyc
     LIBS parameter base device container
-    SOURCES test/test_block_cyclic_distribution.cpp
+    SOURCES test/test_parallel_blkcyc.cpp
+  )
+  AddTest(
+    TARGET MODULE_BASE_basis_flattened_indexing
+    LIBS parameter
+    SOURCES test/test_basis_flattened_indexing.cpp
+            basis_flattened_indexing.cpp
+            parallel_blkcyc.cpp
+            element_basis_index.cpp
   )
+  # TODO: classify the serial and parallel tests
   if(ENABLE_MPI)
     add_test(
-      NAME MODULE_BASE_block_cyclic_distribution_mpi
+      NAME MODULE_BASE_parallel_blkcyc_mpi
       COMMAND
         ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 4
         ${MPIEXEC_PREFLAGS}
-        $<TARGET_FILE:MODULE_BASE_block_cyclic_distribution>
+        $<TARGET_FILE:MODULE_BASE_parallel_blkcyc>
         ${MPIEXEC_POSTFLAGS}
     )
     add_subdirectory(test)

source/source_base/atom_basis_view.cpp

@@ -0,0 +1,77 @@
+#include "source_base/atom_basis_view.h"
+
+#include "source_base/basis_flattened_indexing.h"
+#include "source_base/parallel_blkcyc.h"
+
+#include <stdexcept>
+
+AtomBasisView::AtomBasisView(const Parallel_Blkcyclic* layout,
+                             const int ibas0,
+                             const int local0,
+                             const int nloc,
+                             const bool col_axis)
+    : layout_(layout), ibas0_(ibas0), local0_(local0), nloc_(nloc), col_axis_(col_axis)
+{
+}
+
+AtomBasisView AtomBasisView::make(const BasisFlattenedIndexing& basis,
+                                  const Parallel_Blkcyclic& layout,
+                                  const int iat,
+                                  const bool columns)
+{
+    const int global_size = columns ? layout.ncglb() : layout.nrglb();
+    if (basis.size() != static_cast<std::size_t>(global_size))
+    {
+        throw std::invalid_argument("basis indexing and matrix distribution sizes disagree");
+    }
+
+    const BasisRange basis_range = basis.range(iat);
+    const int ibas0 = static_cast<int>(basis_range.ibas0);
+    int local0 = -1;
+    int last_local = -2;
+    int nloc = 0;
+    for (std::size_t ibas = basis_range.ibas0; ibas < basis_range.ibas1; ++ibas)
+    {
+        const int ibas_global = static_cast<int>(ibas);
+        const int local = columns ? layout.icloc(ibas_global) : layout.irloc(ibas_global);
+        if (local < 0)
+        {
+            continue;
+        }
+        if (nloc == 0)
+        {
+            local0 = local;
+        }
+        else if (local != last_local + 1)
+        {
+            throw std::invalid_argument("atom basis entries are not contiguous");
+        }
+        last_local = local;
+        ++nloc;
+    }
+    return AtomBasisView(&layout, ibas0, local0, nloc, columns);
+}
+
+int AtomBasisView::operator[](const std::size_t index) const
+{
+    // Reproduces the eager value the old vector stored: the global basis index
+    // of the owned local entry minus the atom's first basis index (ibas - ibas0).
+    const int local = local0_ + static_cast<int>(index);
+    const int ibas_global = col_axis_ ? layout_->icglb(local) : layout_->irglb(local);
+    const int entry = ibas_global - ibas0_;
+    return entry;
+}
+
+AtomBasisView AtomBasisView::rows(const BasisFlattenedIndexing& basis,
+                                  const Parallel_Blkcyclic& layout,
+                                  const int iat)
+{
+    return make(basis, layout, iat, false);
+}
+
+AtomBasisView AtomBasisView::cols(const BasisFlattenedIndexing& basis,
+                                  const Parallel_Blkcyclic& layout,
+                                  const int iat)
+{
+    return make(basis, layout, iat, true);
+}

source/source_base/atom_basis_view.h

@@ -0,0 +1,64 @@
+#ifndef ABACUS_SOURCE_BASE_ATOM_BASIS_VIEW_H
+#define ABACUS_SOURCE_BASE_ATOM_BASIS_VIEW_H
+
+#include <cstddef>
+
+class BasisFlattenedIndexing; // Defined in source/source_base/basis_flattened_indexing.h
+class Parallel_Blkcyclic; // Defined in source/source_base/parallel_blkcyc.h
+
+/**
+ * @brief Atom-local basis entries derived from basis indexing and distribution.
+ *
+ * Entries are atom-local basis numbers in local matrix order, materialized on
+ * demand from the indexing and distribution inputs; the view allocates nothing.
+ */
+class AtomBasisView
+{
+  public:
+    static AtomBasisView rows(const BasisFlattenedIndexing& basis,
+                              const Parallel_Blkcyclic& layout,
+                              int iat);
+    static AtomBasisView cols(const BasisFlattenedIndexing& basis,
+                              const Parallel_Blkcyclic& layout,
+                              int iat);
+
+    int nloc() const
+    {
+        return nloc_;
+    }
+
+    bool empty() const
+    {
+        return nloc_ == 0;
+    }
+
+    int local0() const
+    {
+        return local0_;
+    }
+
+    // index must be smaller than nloc(); out-of-range access is not checked.
+    // This view borrows the distribution it was derived from; it must not
+    // outlive it.
+    int operator[](const std::size_t index) const;
+
+  private:
+    AtomBasisView(const Parallel_Blkcyclic* layout,
+                  const int ibas0,
+                  const int local0,
+                  const int nloc,
+                  const bool col_axis);
+
+    static AtomBasisView make(const BasisFlattenedIndexing& basis,
+                              const Parallel_Blkcyclic& layout,
+                              const int iat,
+                              const bool columns);
+
+    const Parallel_Blkcyclic* layout_ = nullptr;
+    int ibas0_ = 0;
+    int local0_ = 0;
+    int nloc_ = 0;
+    bool col_axis_ = false;
+};
+
+#endif