Add SLATE for Distributed Memory Linear Algebra

.gitlab-ci.yml

@@ -7,6 +7,7 @@ default:
 
 variables:
   MAD_NUM_THREADS : 2
+  OMP_NUM_THREADS : 1
   TA_TARGETS : "tiledarray examples-tiledarray ta_test check-tiledarray"
   # Debug builds with ScaLAPACK=ON need increased TA_UT_CTEST_TIMEOUT
   TA_CONFIG : >
@@ -18,6 +19,7 @@ variables:
     ${BLA_VENDOR}
     ${BLA_THREADS}
     ${ENABLE_SCALAPACK}
+    ${ENABLE_SLATE}
 
 before_script:
   # NB: if CMAKE_BUILD_PARALLEL_LEVEL is not set (i.e. using shared runner), use 1 to ensure we have enough memory
@@ -35,6 +37,7 @@ ubuntu:
   variables:
     TA_PYTHON : "TA_PYTHON=ON"
     ENABLE_SCALAPACK : "ENABLE_SCALAPACK=OFF"
+    ENABLE_SLATE     : "ENABLE_SLATE=OFF"
   script:
     - ./ci/.build-project
         --build ./build
@@ -69,6 +72,7 @@ ubuntu:
         CXX: [ g++, clang++-13 ]
         BUILD_TYPE : [ "Release", "Debug" ]
         ENABLE_SCALAPACK : [ "ENABLE_SCALAPACK=ON", "ENABLE_SCALAPACK=OFF" ]
+        ENABLE_SLATE : [ "ENABLE_SLATE=ON", "ENABLE_SLATE=OFF" ]
         RUNNER_TAGS: [ linux ]
       - IMAGE : [ "ubuntu:22.04", "ubuntu:20.04" ]
         CXX: [ g++ ]

CMakeLists.txt

@@ -121,6 +121,9 @@ add_feature_info(MPI ENABLE_MPI "Message-Passing Interface supports distributed-
 option(ENABLE_SCALAPACK "Enable ScaLAPACK Bindings in TiledArray" OFF)
 add_feature_info(ScaLAPACK ENABLE_SCALAPACK "ScaLAPACK provides distributed linear algebra")
 
+option(ENABLE_SLATE "Enable SLATE Bindings in TiledArray" OFF)
+add_feature_info(SLATE ENABLE_SLATE "SLATE provides distributed linear algebra on GPU architectures")
+
 option(ENABLE_WFN91_LINALG_DISCOVERY_KIT "Use linear algebra discovery kit from github.com/wavefunction91 [recommended]" ON)
 add_feature_info(WFN91LinearAlgebraDiscoveryKit ENABLE_WFN91_LINALG_DISCOVERY_KIT "Linear algebra discovery kit from github.com/wavefunction91 supports many more corner cases than the default CMake modules and/or ICL's BLAS++/LAPACK++ modules")
 
@@ -338,6 +341,9 @@ include(${PROJECT_SOURCE_DIR}/cmake/modules/FindOrFetchBoost.cmake)
 if(ENABLE_SCALAPACK)
   include(external/scalapackpp.cmake)
 endif()
+if(ENABLE_SLATE)
+  include(external/slate.cmake)
+endif()
 
 # optional deps:
 # 1. ccache

cmake/modules/FindOrFetchSLATE.cmake

@@ -0,0 +1,32 @@
+# Try find_package
+if (NOT TARGET slate)
+    find_package(slate QUIET CONFIG)
+    if (TARGET slate)
+        message(STATUS "Found SLATE CONFIG at ${slate_CONFIG}")
+    endif (TARGET slate)
+endif (NOT TARGET slate)
+
+# If not found, build via FetchContent
+if (NOT TARGET slate)
+
+    # Make sure BLAS++/LAPACK++ are already in place
+    # (will typically be loaded from BTAS)
+    include(${vg_cmake_kit_SOURCE_DIR}/modules/FindOrFetchLinalgPP.cmake)
+
+    if (NOT TILEDARRAY_HAS_CUDA)
+        set(gpu_backend none CACHE STRING "Device Backend for ICL Linalg++/SLATE")
+    endif (NOT TILEDARRAY_HAS_CUDA)
+
+    include(FetchContent)
+    FetchContent_Declare(
+        slate
+        GIT_REPOSITORY  https://github.com/icl-utk-edu/slate.git
+        GIT_TAG         ${TA_TRACKED_SLATE_TAG}
+    )
+    FetchContent_MakeAvailable(slate)
+
+endif (NOT TARGET slate)
+
+if (NOT TARGET slate)
+    message( FATAL_ERROR "FindOrFetchSLATE could not make slate target available")
+endif (NOT TARGET slate)

examples/CMakeLists.txt

@@ -33,6 +33,7 @@ add_subdirectory (cuda)
 add_subdirectory (dgemm)
 add_subdirectory (demo)
 add_subdirectory (scalapack)
+add_subdirectory (slate)
 add_subdirectory (fock)
 add_subdirectory (mpi_tests)
 add_subdirectory (pmap_test)

examples/slate/CMakeLists.txt

@@ -0,0 +1,38 @@
+#
+#  This file is a part of TiledArray.
+#  Copyright (C) 2016  Virginia Tech
+#
+#  This program is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+#  Drew Lewis
+#  Department of Chemistry, Virginia Tech
+#
+#  CMakeLists.txt
+#  Dec 6th, 2016
+#
+
+# Create example executable
+
+if(ENABLE_SLATE)
+
+foreach(_exec conversion)
+
+  # Add executable
+  add_executable(slate-${_exec} EXCLUDE_FROM_ALL ${_exec}.cpp)
+  target_link_libraries(slate-${_exec} PRIVATE tiledarray)
+  add_dependencies(examples-tiledarray slate-${_exec})
+
+endforeach()
+
+endif(ENABLE_SLATE)

examples/slate/conversion.cpp

@@ -0,0 +1,140 @@
+/*
+ *  This file is a part of TiledArray.
+ *  Copyright (C) 2020  Virginia Tech
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *  David Williams-Young
+ *  Computational Research Division, Lawrence Berkeley National Laboratory
+ *
+ *  conversion.cpp
+ *  Created: 7  Feb, 2020
+ *  Edited:  13 May, 2020
+ *
+ */
+
+#include <TiledArray/conversions/slate.h>
+
+#include <tiledarray.h>
+#include <random>
+#include <slate/slate.hh>
+#include <TiledArray/pmap/user_pmap.h>
+
+
+template <typename Integral1, typename Integral2>
+int64_t div_ceil(Integral1 x, Integral2 y) {
+  int64_t x_ll = x;
+  int64_t y_ll = y;
+
+  auto d = std::div(x_ll, y_ll);
+  return d.quot + !!d.rem;
+}
+
+TA::TiledRange gen_trange(size_t N, const std::vector<size_t>& TA_NBs) {
+  assert(TA_NBs.size() > 0);
+
+  std::default_random_engine gen(0);
+  std::uniform_int_distribution<> dist(0, TA_NBs.size() - 1);
+  auto rand_indx = [&]() { return dist(gen); };
+  auto rand_nb = [&]() { return TA_NBs[rand_indx()]; };
+
+  std::vector<size_t> t_boundaries = {0};
+  auto TA_NB = rand_nb();
+  while (t_boundaries.back() + TA_NB < N) {
+    t_boundaries.emplace_back(t_boundaries.back() + TA_NB);
+    TA_NB = rand_nb();
+  }
+  t_boundaries.emplace_back(N);
+
+  std::vector<TA::TiledRange1> ranges(
+      2, TA::TiledRange1(t_boundaries.begin(), t_boundaries.end()));
+
+  return TA::TiledRange(ranges.begin(), ranges.end());
+};
+
+
+
+auto make_square_proc_grid(MPI_Comm comm) {
+    int mpi_size; MPI_Comm_size(comm, &mpi_size);
+    int p,q;
+    for(p = int( sqrt( mpi_size ) ); p > 0; --p) {
+        q = int( mpi_size / p );
+        if(p*q == mpi_size) break;
+    }
+    return std::make_pair(p,q);
+}
+
+
+
+
+
+
+
+
+
+int main(int argc, char** argv) {
+  auto& world = TA::initialize(argc, argv);
+  {
+    int64_t N = argc > 1 ? std::stoi(argv[1]) : 1000;
+    size_t NB = argc > 2 ? std::stoi(argv[2]) : 128;
+
+    auto make_ta_reference = 
+      [&](TA::Tensor<double>& t, TA::Range const& range) {
+
+        t = TA::Tensor<double>(range, 0.0);
+        auto lo = range.lobound_data();
+        auto up = range.upbound_data();
+        for (int m = lo[0]; m < up[0]; ++m) {
+          for (int n = lo[1]; n < up[1]; ++n) {
+            t(m, n) = m - n;
+          }
+        }
+
+        return t.norm();
+      };
+
+    // Generate Reference TA tensor.
+    auto trange = gen_trange(N, {NB});
+    auto ref_ta =
+        TA::make_array<TA::TArray<double> >(world, trange, make_ta_reference);
+
+    // Do Conversion 
+    auto A = TA::array_to_slate( ref_ta );
+    auto A_ta = TA::slate_to_array<TA::TArray<double>>(A, world);
+    world.gop.fence();
+
+    // Slate matrix to eigen
+    Eigen::MatrixXd slate_eigen = Eigen::MatrixXd::Zero(N,N);
+    for (int64_t j = 0; j < A.nt(); ++j) 
+    for (int64_t i = 0; i < A.mt(); ++i) {
+        A.tileBcast(i,j, A, slate::Layout::ColMajor);
+        auto T = A(i,j);
+        Eigen::Map<Eigen::MatrixXd> T_map( T.data(), T.mb(), T.nb() );
+        slate_eigen.block(i*NB,j*NB,T.mb(), T.nb()) = T_map; 
+    }
+    //if(!world.rank()) {
+    //std::cout << "SLATE\n" << slate_eigen << std::endl;
+    //}
+
+    A_ta.make_replicated();
+    world.gop.fence();
+    auto A_eigen = TA::array_to_eigen(A_ta);
+    //if(!world.rank()) std::cout << "TA\n" << A_eigen << std::endl;
+    std::cout << (A_eigen - slate_eigen).norm() << std::endl;
+
+
+  }
+
+  TA::finalize();
+}

external/slate.cmake

@@ -0,0 +1,21 @@
+if (NOT TARGET slate)
+  set(VGCMAKEKIT_TRACKED_SLATE_TAG ${TA_TRACKED_SLATE_TAG} CACHE STRING "slate tag")
+  include(FindOrFetchSLATE)
+endif()
+
+# built {blacs,scalapack}pp as a subproject? install as part of tiledarray export as well
+# to be able to use TiledArray_SLATE from the build tree
+if (TARGET slate)
+    install( TARGETS slate EXPORT tiledarray COMPONENT tiledarray )
+    # Add these dependencies to External
+    add_dependencies(External-tiledarray slate )
+endif()
+
+if (TARGET slate)
+  add_library( TiledArray_SLATE INTERFACE )
+  target_link_libraries( TiledArray_SLATE INTERFACE slate )
+
+  install( TARGETS TiledArray_SLATE EXPORT tiledarray COMPONENT tiledarray )
+
+  set( TILEDARRAY_HAS_SLATE 1 )
+endif()

external/versions.cmake

@@ -36,6 +36,9 @@ set(TA_TRACKED_UMPIRE_PREVIOUS_TAG v6.0.0)
 set(TA_TRACKED_SCALAPACKPP_TAG 6397f52cf11c0dfd82a79698ee198a2fce515d81)
 set(TA_TRACKED_SCALAPACKPP_PREVIOUS_TAG 711ef363479a90c88788036f9c6c8adb70736cbf )
 
+set(TA_TRACKED_SLATE_TAG 04d552f12e0d181c27652ecd3ebedb265b7eec07)
+set(TA_TRACKED_SLATE_PREVIOUS_TAG 8651441aa87cd69b560d4dac8c5ceb0e7f8c32a4)
+
 set(TA_TRACKED_RANGEV3_TAG 2e0591c57fce2aca6073ad6e4fdc50d841827864)
 set(TA_TRACKED_RANGEV3_PREVIOUS_TAG dbdaa247a25a0daa24c68f1286a5693c72ea0006)
 

src/CMakeLists.txt

@@ -274,6 +274,9 @@ endif(CUDA_FOUND)
 if( TARGET TiledArray_SCALAPACK )
   list(APPEND _TILEDARRAY_DEPENDENCIES TiledArray_SCALAPACK)
 endif()
+if( TARGET TiledArray_SLATE )
+  list(APPEND _TILEDARRAY_DEPENDENCIES TiledArray_SLATE)
+endif()
 list(APPEND _TILEDARRAY_DEPENDENCIES "${LAPACK_LIBRARIES}")
 
 if( TARGET ttg-parsec )

src/TiledArray/config.h.in

@@ -74,6 +74,9 @@
 /* Define if TA has enabled ScaLAPACK Bindings */
 #cmakedefine TILEDARRAY_HAS_SCALAPACK 1
 
+/* Define if TA has enabled SLATE Bindings */
+#cmakedefine TILEDARRAY_HAS_SLATE 1
+
 /* Define if TiledArray configured with CUDA support */
 #cmakedefine TILEDARRAY_HAS_CUDA @TILEDARRAY_HAS_CUDA@
 #cmakedefine TILEDARRAY_CHECK_CUDA_ERROR @TILEDARRAY_CHECK_CUDA_ERROR@