cuTENSOR Tensor contraction

.github/.licenserc.yaml

@@ -22,6 +22,7 @@ header:
     - docs/Makefile
     - LICENSE
     - cmake/config.hpp.in
+    - cmake/FindcuTENSOR.cmake
     - docs/requirements.txt
     - docs/source/bibliography/*.bib
     - version.txt

CMakeLists.txt

@@ -29,16 +29,34 @@ include(get_cmaize)
 set(project_inc_dir "${CMAKE_CURRENT_LIST_DIR}/include/${PROJECT_NAME}")
 set(project_src_dir "${CMAKE_CURRENT_LIST_DIR}/src/${PROJECT_NAME}")
 
+# Documentation
 include(nwx_cxx_api_docs)
 nwx_cxx_api_docs("${project_inc_dir}" "${project_src_dir}")
 
+## Extensions ##
+set(SOURCE_EXTS "cpp")
+set(INCLUDE_EXTS "hpp")
+
 ### Options ###
 cmaize_option_list(
     BUILD_TESTING OFF "Should we build the tests?"
     BUILD_PYBIND11_PYBINDINGS ON "Should we build Python3 bindings?"
     ENABLE_SIGMA OFF "Should we enable Sigma for uncertainty tracking?"
+    ENABLE_CUTENSOR OFF "Should we enable cuTENSOR?"
 )
 
+if("${ENABLE_CUTENSOR}") 
+    if("${ENABLE_SIGMA}")
+        set(MSG "Sigma is not compatible with cuTENSOR. Turning Sigma OFF.")
+        message(WARNING ${MSG})
+        set(ENABLE_SIGMA OFF)
+    endif()
+    enable_language(CUDA)
+    set(SOURCE_EXTS ${SOURCE_EXTS} cu)
+    set(INCLUDE_EXTS ${INCLUDE_EXTS} cuh hu)
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
+endif()
+
 ### Dependendencies ###
 include(get_utilities)
 
@@ -64,15 +82,27 @@ cmaize_find_or_build_optional_dependency(
     CMAKE_ARGS BUILD_TESTING=OFF
                ENABLE_EIGEN_SUPPORT=ON
 )
+set(DEPENDENCIES utilities parallelzone Boost::boost eigen sigma)
+
+if("${ENABLE_CUTENSOR}")
+    include(cmake/FindcuTENSOR.cmake)
+    list(APPEND DEPENDENCIES cuTENSOR::cuTENSOR)
+endif()
 
 cmaize_add_library(
     ${PROJECT_NAME}
     SOURCE_DIR "${project_src_dir}"
+    SOURCE_EXTS "${SOURCE_EXTS}"
     INCLUDE_DIRS "${project_inc_dir}"
-    DEPENDS utilities parallelzone Boost::boost eigen sigma
+    INCLUDE_EXTS "${INCLUDE_EXTS}"
+    DEPENDS "${DEPENDENCIES}"
 )
 target_include_directories(${PROJECT_NAME} PUBLIC "${CMAKE_CURRENT_BINARY_DIR}")
 
+if("${ENABLE_CUTENSOR}")
+    target_compile_definitions("${PROJECT_NAME}" PUBLIC ENABLE_CUTENSOR)
+endif()
+
 include(nwx_pybind11)
 nwx_add_pybind11_module(
     ${PROJECT_NAME}

cmake/FindcuTENSOR.cmake

@@ -0,0 +1,135 @@
+#=============================================================================
+# Copyright (c) 2021, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#=============================================================================
+
+#[=======================================================================[.rst:
+FindcuTENSOR
+--------
+
+Find cuTENSOR
+
+Imported targets
+^^^^^^^^^^^^^^^^
+
+This module defines the following :prop_tgt:`IMPORTED` target(s):
+
+``cuTENSOR::cuTENSOR``
+  The cuTENSOR library, if found.
+
+Result variables
+^^^^^^^^^^^^^^^^
+
+This module will set the following variables in your project:
+
+``cuTENSOR_FOUND``
+  True if cuTENSOR is found.
+``cuTENSOR_INCLUDE_DIRS``
+  The include directories needed to use cuTENSOR.
+``cuTENSOR_LIBRARIES``
+  The libraries needed to usecuTENSOR.
+``cuTENSOR_VERSION_STRING``
+  The version of the cuTENSOR library found. [OPTIONAL]
+
+#]=======================================================================]
+
+# Prefer using a Config module if it exists for this project
+set(cuTENSOR_NO_CONFIG FALSE)
+if(NOT cuTENSOR_NO_CONFIG)
+  find_package(cuTENSOR CONFIG QUIET HINTS ${cutensor_DIR})
+  if(cuTENSOR_FOUND)
+    find_package_handle_standard_args(cuTENSOR DEFAULT_MSG cuTENSOR_CONFIG)
+    return()
+  endif()
+endif()
+
+find_path(cuTENSOR_INCLUDE_DIR NAMES cutensor.h )
+
+set(cuTENSOR_IS_HEADER_ONLY FALSE)
+if(NOT cuTENSOR_LIBRARY AND NOT cuTENSOR_IS_HEADER_ONLY)
+  find_library(cuTENSOR_LIBRARY_RELEASE NAMES libcutensor.so NAMES_PER_DIR )
+  find_library(cuTENSOR_LIBRARY_DEBUG   NAMES libcutensor.sod   NAMES_PER_DIR )
+
+  include(${CMAKE_ROOT}/Modules/SelectLibraryConfigurations.cmake)
+  select_library_configurations(cuTENSOR)
+  unset(cuTENSOR_FOUND) #incorrectly set by select_library_configurations
+endif()
+
+include(${CMAKE_ROOT}/Modules/FindPackageHandleStandardArgs.cmake)
+
+if(cuTENSOR_IS_HEADER_ONLY)
+  find_package_handle_standard_args(cuTENSOR
+                                    REQUIRED_VARS cuTENSOR_INCLUDE_DIR
+                                    VERSION_VAR )
+else()
+  find_package_handle_standard_args(cuTENSOR
+                                    REQUIRED_VARS cuTENSOR_LIBRARY cuTENSOR_INCLUDE_DIR
+                                    VERSION_VAR )
+endif()
+
+if(NOT cuTENSOR_FOUND)
+  set(CUTENSOR_FILENAME libcutensor-linux-x86_64-${CUTENSOR_VERSION}-archive)
+
+  message(STATUS "cuTENSOR not found. Downloading library. By continuing this download you accept to the license terms of cuTENSOR")
+
+  CPMAddPackage(
+    NAME cutensor
+    VERSION ${CUTENSOR_VERSION}
+    URL https://developer.download.nvidia.com/compute/cutensor/redist/libcutensor/linux-x86_64/libcutensor-linux-x86_64-${CUTENSOR_VERSION}-archive.tar.xz
+    # Eigen's CMakelists are not intended for library use
+    DOWNLOAD_ONLY YES 
+  )
+
+  set(cuTENSOR_LIBRARY ${cutensor_SOURCE_DIR}/lib/${CUDAToolkit_VERSION_MAJOR}/libcutensor.so)
+  set(cuTENSOR_INCLUDE_DIR ${cutensor_SOURCE_DIR}/include) 
+
+
+  set(cuTENSOR_FOUND TRUE)
+endif()
+
+if(cuTENSOR_FOUND)
+  set(cuTENSOR_INCLUDE_DIRS ${cuTENSOR_INCLUDE_DIR})
+
+  if(NOT cuTENSOR_LIBRARIES)
+    set(cuTENSOR_LIBRARIES ${cuTENSOR_LIBRARY})
+  endif()
+
+  if(NOT TARGET cuTENSOR::cuTENSOR)
+    add_library(cuTENSOR::cuTENSOR UNKNOWN IMPORTED)
+    set_target_properties(cuTENSOR::cuTENSOR PROPERTIES
+      INTERFACE_INCLUDE_DIRECTORIES "${cuTENSOR_INCLUDE_DIRS}")
+
+    if(cuTENSOR_LIBRARY_RELEASE)
+      set_property(TARGET cuTENSOR::cuTENSOR APPEND PROPERTY
+        IMPORTED_CONFIGURATIONS RELEASE)
+      set_target_properties(cuTENSOR::cuTENSOR PROPERTIES
+        IMPORTED_LOCATION_RELEASE "${cuTENSOR_LIBRARY_RELEASE}")
+    endif()
+
+    if(cuTENSOR_LIBRARY_DEBUG)
+      set_property(TARGET cuTENSOR::cuTENSOR APPEND PROPERTY
+        IMPORTED_CONFIGURATIONS DEBUG)
+      set_target_properties(cuTENSOR::cuTENSOR PROPERTIES
+        IMPORTED_LOCATION_DEBUG "${cuTENSOR_LIBRARY_DEBUG}")
+    endif()
+
+    if(NOT cuTENSOR_LIBRARY_RELEASE AND NOT cuTENSOR_LIBRARY_DEBUG)
+      set_property(TARGET cuTENSOR::cuTENSOR APPEND PROPERTY
+        IMPORTED_LOCATION "${cuTENSOR_LIBRARY}")
+    endif()
+  endif()
+endif()
+
+unset(cuTENSOR_NO_CONFIG)
+unset(cuTENSOR_IS_HEADER_ONLY)

src/tensorwrapper/buffer/detail_/cutensor_traits.cuh

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+#ifdef ENABLE_CUTENSOR
+#include <cuda_runtime.h>
+#include <cutensor.h>
+
+namespace tensorwrapper::buffer::detail_ {
+
+// Traits for cuTENSOR based on the floating point type
+template<typename FloatType>
+struct cutensor_traits {};
+
+template<>
+struct cutensor_traits<float> {
+    cutensorDataType_t cutensorDataType     = CUTENSOR_R_32F;
+    cutensorComputeDescriptor_t descCompute = CUTENSOR_COMPUTE_DESC_32F;
+};
+
+template<>
+struct cutensor_traits<double> {
+    cutensorDataType_t cutensorDataType     = CUTENSOR_R_64F;
+    cutensorComputeDescriptor_t descCompute = CUTENSOR_COMPUTE_DESC_64F;
+};
+
+} // namespace tensorwrapper::buffer::detail_
+
+#endif

src/tensorwrapper/buffer/detail_/eigen_tensor.cpp

@@ -18,6 +18,10 @@
 #include "../contraction_planner.hpp"
 #include "eigen_tensor.hpp"
 
+#ifdef ENABLE_CUTENSOR
+#include "eigen_tensor.cuh"
+#endif
+
 namespace tensorwrapper::buffer::detail_ {
 
 #define TPARAMS template<typename FloatType, unsigned int Rank>
@@ -96,6 +100,16 @@ void EIGEN_TENSOR::contraction_assignment_(label_type olabels,
                                            const_pimpl_reference rhs) {
     ContractionPlanner plan(olabels, llabels, rlabels);
 
+#ifdef ENABLE_CUTENSOR
+    // Prepare m_tensor_
+    m_tensor_ = allocate_from_shape_(result_shape.as_smooth(),
+                                     std::make_index_sequence<Rank>());
+    m_tensor_.setZero();
+
+    // Dispatch to cuTENSOR
+    cutensor_contraction<my_type>(olabels, llabels, rlabels, result_shape, lhs,
+                                  rhs, m_tensor_);
+#else
     auto lt = lhs.clone();
     auto rt = rhs.clone();
     lt->permute_assignment(plan.lhs_permutation(), llabels, lhs);
@@ -140,6 +154,7 @@ void EIGEN_TENSOR::contraction_assignment_(label_type olabels,
     } else {
         m_tensor_ = tensor;
     }
+#endif
     mark_for_rehash_();
 }