Tensor transpose CPU+GPU

CMakeLists.txt

@@ -37,7 +37,7 @@ endif()
 
 set(CMSB_PROJECTS TAMM)
 
-set(TAMM_DEPENDENCIES NJSON DOCTEST SPDLOG Eigen3 HPTT)
+set(TAMM_DEPENDENCIES NJSON DOCTEST SPDLOG Eigen3)
 
 if(NOT DEFINED USE_HDF5 OR USE_HDF5)
   list(APPEND TAMM_DEPENDENCIES HDF5)
@@ -50,14 +50,6 @@ elseif(NOT DEFINED USE_LIBNUMA OR USE_LIBNUMA)
   list(APPEND TAMM_DEPENDENCIES numactl)
 endif()
 
-if(TAMM_ENABLE_CUDA OR TAMM_ENABLE_HIP OR TAMM_ENABLE_DPCPP)
-  if(USE_SYCL_BLAS)
-    #list(APPEND TAMM_DEPENDENCIES sycl-blas)
-  endif()
-
-  list(APPEND TAMM_DEPENDENCIES Librett)
-endif()
-
 if(TAMM_ENABLE_ELPA)
   list(APPEND TAMM_DEPENDENCIES ELPA)
 endif()

docs/install.rst

@@ -1,5 +1,3 @@
-The prerequisites needed to build this repository can be found :doc:`here <prerequisites>`.
-
 Build Instructions
 ==================
 
@@ -10,14 +8,13 @@ Dependencies
 
 * cmake >= 3.26
 * MPI 
-* C++17 compiler (information on supported compilers here :doc:`here <prerequisites>`.)
+* C++20 Compilers (GNU >= 14.1, LLVM Clang >= 18, AppleClang >= 19)
 * CUDA >= 12.8 (Required only for CUDA builds)
 * ROCM >= 6.2  (Required only for ROCM builds)
 
 **The remaining dependencies are automatically built and do not need to be installed explicitly:**
 
 * GlobalArrays
-* HPTT, Librett
 * HDF5
 * BLAS/LAPACK (BLIS and netlib-lapack are automatically built if no vendor BLAS libraries are provided)
 * BLAS++ and LAPACK++
@@ -124,7 +121,8 @@ Default build on MACOS
 ~~~~~~~~~~~~~~~~~~~~~~
 
 .. note::
-   The prerequisites for ``MACOS`` can be installed using ``brew`` as detailed :doc:`here <prerequisites>`.
+   The required external dependencies for ``MACOS`` can be installed using ``brew``:
+   ``brew install gcc mpich cmake autoconf automake libtool``
 
 ::
 
@@ -203,7 +201,7 @@ Polaris modules and env
 
    module use /soft/modulefiles/
    module load PrgEnv-gnu
-   module load cudatoolkit-standalone/12.6.1 spack-pe-base cmake
+   module load cudatoolkit-standalone spack-pe-base cmake
    module unload craype-accel-nvidia80
 
    export CRAYPE_LINK_TYPE=dynamic

src/tamm/CMakeLists.txt

@@ -10,7 +10,18 @@ if(${PROJECT_NAME}_HAS_CUDA OR ${PROJECT_NAME}_HAS_HIP OR ${PROJECT_NAME}_HAS_DP
     mr/per_device_resource.hpp
   )
 
-  set(tamm_gpu_SRCS kernels/gpu_blas.cpp kernels/gpu_lapack.cpp)
+  set(tamm_gpu_SRCS
+    kernels/gpu_blas.cpp
+    kernels/gpu_lapack.cpp
+    kernels/gpu_permute.cpp
+  )
+  if(${PROJECT_NAME}_HAS_CUDA)
+    # gpu_permute.cpp holds __global__ kernels + <<<>>> launches, so nvcc
+    # (not the C++ compiler) must compile it even though it is a .cpp file.
+    set_source_files_properties(kernels/gpu_permute.cpp PROPERTIES LANGUAGE CUDA)
+  endif()
+  # NOTE: hipcc compiles .cpp HIP code directly, and icpx compiles the SYCL
+  # branch, so no LANGUAGE override is needed for HIP/DPCPP builds.
 endif()
 if(${PROJECT_NAME}_HAS_DPCPP)
   list(APPEND tamm_gpu_INCS sycl_device.hpp)
@@ -74,6 +85,8 @@ set(TAMM_INCLUDES
     kernels/assign.hpp
     kernels/multiply.hpp
     kernels/tamm_blas.hpp
+    kernels/gpu_permute.hpp
+    kernels/cpu_permute.hpp
     op_dag.hpp
     interfaces.hpp
     op_profiler.hpp
@@ -90,7 +103,7 @@ set(TAMM_INCLUDES
     blockops_cpu.hpp
     perm.hpp
     iteration.hpp
-    ip_hptt.hpp
+    index_permute.hpp
     blockops_blas.hpp
     tensor_variant.hpp
     op_executor.hpp

src/tamm/block_assign_plan.hpp

@@ -6,7 +6,7 @@
 #include "tamm/blockops_blas.hpp"
 #include "tamm/blockops_cpu.hpp"
 #include "tamm/errors.hpp"
-#include "tamm/ip_hptt.hpp"
+#include "tamm/index_permute.hpp"
 #include "tamm/types.hpp"
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -19,9 +19,6 @@ namespace tamm {
 /**
  * @brief Block assign plan for block assignment on Tensor Ops
  *
- * @todo HPTT calls should be fixed for mixed precision (T1 && T2) when we allow
- * those operations in make_op, currently only complex and scalar mix is allowed
- *
  */
 class BlockAssignPlan {
 public:
@@ -37,7 +34,7 @@ class BlockAssignPlan {
     // rhs_labels{rhs_labels},
     plan_{Plan::invalid} {
     prep_flat_plan(lhs_labels, rhs_labels);
-    if(plan_ == Plan::invalid) { prep_hptt(lhs_labels, rhs_labels); }
+    if(plan_ == Plan::invalid) { prep_permute(lhs_labels, rhs_labels); }
     if(plan_ == Plan::invalid) { prep_index_permute_plan(lhs_labels, rhs_labels); }
     if(plan_ == Plan::invalid) { prep_general_loop_plan(lhs_labels, rhs_labels); }
     if(plan_ == Plan::invalid) { prep_general_plan(lhs_labels, rhs_labels); }
@@ -56,9 +53,9 @@ class BlockAssignPlan {
     switch(plan_) {
       case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rhs); break;
       case Plan::flat_update: blockops::cpu::flat_update(lhs, rhs); break;
-      case Plan::hptt:
-        blockops::hptt::index_permute_hptt(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_,
-                                           rhs.block_dims());
+      case Plan::permute:
+        blockops::permute::index_permute(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_,
+                                                 rhs.block_dims());
         break;
       case Plan::index_permute_assign:
         blockops::cpu::index_permute_assign(lhs.buf(), rhs.buf(), ip_plan_.perm_, lhs.block_dims());
@@ -118,9 +115,9 @@ class BlockAssignPlan {
     switch(plan_) {
       case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rscale, rhs); break;
       case Plan::flat_update: blockops::cpu::flat_update(lhs, rscale, rhs); break;
-      case Plan::hptt:
-        blockops::hptt::index_permute_hptt(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_,
-                                           rhs.block_dims());
+      case Plan::permute:
+        blockops::permute::index_permute(lscale, lhs.buf(), rscale, rhs.buf(),
+                                                 ip_plan_.perm_, rhs.block_dims());
         break;
       case Plan::index_permute_assign:
         blockops::cpu::index_permute_assign(lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_,
@@ -200,9 +197,9 @@ class BlockAssignPlan {
     switch(plan_) {
       case Plan::flat_assign: NOT_ALLOWED(); break;
       case Plan::flat_update: blockops::cpu::flat_update(lscale, lhs, rscale, rhs); break;
-      case Plan::hptt:
-        blockops::hptt::index_permute_hptt(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_,
-                                           rhs.block_dims());
+      case Plan::permute:
+        blockops::permute::index_permute(lscale, lhs.buf(), rscale, rhs.buf(),
+                                                 ip_plan_.perm_, rhs.block_dims());
         break;
       case Plan::index_permute_assign: NOT_ALLOWED(); break;
       case Plan::index_permute_update:
@@ -268,12 +265,12 @@ class BlockAssignPlan {
   }
 
   template<typename T>
-  void prep_hptt(const std::vector<T>& lhs_labels, const std::vector<T>& rhs_labels) {
+  void prep_permute(const std::vector<T>& lhs_labels, const std::vector<T>& rhs_labels) {
     prep_index_permute_plan(lhs_labels, rhs_labels);
     if(plan_ == Plan::invalid) {
-      return; // no index permute plan, no hptt plan
+      return; // no index permute plan, no permute plan
     }
-    plan_ = Plan::hptt;
+    plan_ = Plan::permute;
   }
 
   template<typename T>
@@ -318,7 +315,7 @@ class BlockAssignPlan {
   enum class Plan {
     flat_assign,
     flat_update,
-    hptt,
+    permute,
     index_permute_assign,
     index_permute_update,
     ipgen_loop_assign,
@@ -346,7 +343,7 @@ class BlockAssignPlan {
     switch(plan) {
       case Plan::flat_assign: return "flat_assign";
       case Plan::flat_update: return "flat_update";
-      case Plan::hptt: return "hptt";
+      case Plan::permute: return "permute";
       case Plan::index_permute_assign: return "index_permute_assign";
       case Plan::index_permute_update: return "index_permute_update";
       case Plan::ipgen_loop_assign: return "ipgen_loop_assign";

src/tamm/index_permute.hpp

@@ -0,0 +1,43 @@
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+#include "tamm/errors.hpp"
+#include "tamm/kernels/cpu_permute.hpp"
+#include "tamm/types.hpp"
+
+namespace tamm::blockops::permute {
+
+///////////////////////////////////////////////////////////////////////////////
+//
+//           Index permute using the in-house CPU permute kernel
+//
+///////////////////////////////////////////////////////////////////////////////
+
+// Out-of-place permute with assign/accumulate semantics:
+//   lbuf = lscale * lbuf + rscale * permute(rbuf)
+// perm_to_dest maps each destination axis to its source axis and sdims holds
+// the source extents, both in natural order; the output extents follow as
+// outDims[k] == sdims[perm_to_dest[k]]. Scale types are independent template
+// parameters so call sites may mix e.g. an int literal rscale with a
+// complex buffer.
+template<typename BL, typename T1, typename BA, typename T2>
+void index_permute(BL lscale, T1* lbuf, BA rscale, const T2* rbuf,
+                   const PermVector& perm_to_dest, const std::vector<size_t>& sdims) {
+  const size_t ndim = sdims.size();
+  EXPECTS(perm_to_dest.size() == ndim);
+  EXPECTS(ndim <= static_cast<size_t>(kernels::gpu::permute_maxrank));
+  EXPECTS(lbuf != nullptr && (ndim == 0 || rbuf != nullptr));
+
+  size_t outDims[kernels::gpu::permute_maxrank] = {};
+  int    perm[kernels::gpu::permute_maxrank]    = {};
+  for(size_t i = 0; i < ndim; ++i) {
+    perm[i] = static_cast<int>(perm_to_dest[i]);
+    EXPECTS(perm[i] >= 0 && static_cast<size_t>(perm[i]) < ndim);
+    outDims[i] = sdims[static_cast<size_t>(perm[i])];
+  }
+  kernels::cpu::permute(lbuf, rbuf, static_cast<int>(ndim), outDims, perm, rscale, lscale);
+}
+
+} // namespace tamm::blockops::permute