improve cutensor lib discovery inc. conda install, allow custom tblis…

.github/workflows/cmake.yml

@@ -32,10 +32,12 @@ jobs:
           - os: ubuntu-24.04
             cc: /usr/bin/gcc-14
             cxx: /usr/bin/g++-14
+            cuda: true
             sanitize_flags: -fsanitize=address -fsanitize=leak -fsanitize=undefined -fno-omit-frame-pointer -fno-var-tracking
           - os: macos-14
             cc: clang
             cxx: clang++
+            cuda: false
             sanitize_flags: -fsanitize=address -fsanitize=undefined -fno-omit-frame-pointer -fno-var-tracking
 
     name: "${{ matrix.valgrind && 'Valgrind' || matrix.sanitize && 'Sanitizers' || '' }} ${{ matrix.os }}: ${{ matrix.cxx }} ${{ matrix.build_type }}"
@@ -52,6 +54,8 @@ jobs:
         -DCMAKE_BUILD_TYPE=${{ matrix.build_type }}
         -DCMAKE_UNITY_BUILD=${{ matrix.build_type == 'Debug' || matrix.valgrind }}
         -DTAPP_REFERENCE_ENABLE_TBLIS=ON
+        -DTAPP_REFERENCE_BUILD_CUTENSOR_BINDS=${{ matrix.cuda && 'ON' || 'OFF' }}
+
     steps:
     - uses: actions/checkout@v4
 
@@ -98,6 +102,17 @@ jobs:
       run: |
         sudo apt-get update
         sudo apt-get install ninja-build g++-14 liblapack-dev ccache valgrind
+
+    - name: Install prerequisites CUDA Toolkit (Ubuntu only)
+      if: ${{ matrix.cuda }}
+      run: |
+        sudo apt-get install -y nvidia-cuda-toolkit
+
+    - name: Set CUDA host compiler
+      if: ${{ matrix.cuda }}
+      run: |
+        echo "CUDAHOSTCXX=${{ matrix.cxx }}" >> $GITHUB_ENV
+
     - name: Prepare ccache timestamp
       id: ccache_cache_timestamp
       shell: cmake -P {0}

CMakeLists.txt

@@ -114,18 +114,30 @@ if(TAPP_REFERENCE_ENABLE_TBLIS)
   endif()
 
   set(TBLIS_INSTALL_DIR ${CMAKE_CURRENT_BINARY_DIR}/tblis)
+  # Option to provide custom path to tblis source
+  set(TAPP_REFERENCE_TBLIS_SOURCE_DIR "" CACHE PATH "Path to existing tblis source directory (if empty, will default to fetching from GitHub)")
+
+  if(TAPP_REFERENCE_TBLIS_SOURCE_DIR)
+    # Use user-provided tblis source directory
+    if(NOT EXISTS "${TAPP_REFERENCE_TBLIS_SOURCE_DIR}/CMakeLists.txt")
+      message(FATAL_ERROR "TAPP_REFERENCE_TBLIS_SOURCE_DIR is set to '${TAPP_REFERENCE_TBLIS_SOURCE_DIR}' but no CMakeLists.txt found there")
+    endif()
+    message(STATUS "Using tblis from: ${TAPP_REFERENCE_TBLIS_SOURCE_DIR}")
+    add_subdirectory(${TAPP_REFERENCE_TBLIS_SOURCE_DIR} ${CMAKE_CURRENT_BINARY_DIR}/_deps/tblis-build)
+  else()
+    # Fetch tblis from GitHub
+    include(FetchContent)
+
+    FetchContent_Declare(
+      tblis
+      GIT_REPOSITORY https://github.com/devinamatthews/tblis.git
+      GIT_TAG 9b95712
+      PREFIX ${CMAKE_CURRENT_BINARY_DIR}/_deps/tblis
+      UPDATE_DISCONNECTED TRUE
+    )
 
-  include(FetchContent)
-
-  FetchContent_Declare(
-    tblis
-    GIT_REPOSITORY https://github.com/devinamatthews/tblis.git
-    GIT_TAG 9b95712
-    PREFIX ${CMAKE_CURRENT_BINARY_DIR}/_deps/tblis
-    UPDATE_DISCONNECTED TRUE
-  )
-
-  FetchContent_MakeAvailable(tblis)
+    FetchContent_MakeAvailable(tblis)
+  endif()
 
   target_compile_definitions(tapp-reference PRIVATE TAPP_REFERENCE_ENABLE_TBLIS=1)
 
@@ -175,7 +187,8 @@ if (TAPP_REFERENCE_BUILD_CUTENSOR_BINDINGS)
 
   set(CUTENSOR_ROOT "/usr/local/cutensor")
   set(CUTENSOR_INCLUDE_DIR "${CUTENSOR_ROOT}/include")
-  set(CUTENSOR_LIBRARY_DIR "${CUTENSOR_ROOT}/lib" "${CUTENSOR_ROOT}/lib/11")
+  file(GLOB CUTENSOR_VERSIONED_DIRS "${CUTENSOR_ROOT}/lib/[0-9]*")
+  set(CUTENSOR_LIBRARY_DIR "${CUTENSOR_ROOT}/lib" ${CUTENSOR_VERSIONED_DIRS})
 
   find_library(
     CUTENSOR_LIB
@@ -185,9 +198,18 @@ if (TAPP_REFERENCE_BUILD_CUTENSOR_BINDINGS)
 
   if (NOT CUTENSOR_LIB)
     message(FATAL_ERROR "cuTENSOR library not found. Set CUTENSOR_ROOT correctly.")
+  else()
+    get_filename_component(CUTENSOR_LIBRARY_DIR ${CUTENSOR_LIB} DIRECTORY)
+    if(CUTENSOR_LIBRARY_DIR MATCHES "/[0-9]+$")
+      get_filename_component(CUTENSOR_LIBRARY_DIR ${CUTENSOR_LIBRARY_DIR} DIRECTORY)
+    endif()
+    get_filename_component(CUTENSOR_ROOT ${CUTENSOR_LIBRARY_DIR} DIRECTORY)
+
+    set(CUTENSOR_INCLUDE_DIR "${CUTENSOR_ROOT}/include")
   endif()
 
   message(STATUS "Found cuTENSOR: ${CUTENSOR_LIB}")
+  message(STATUS "cuTENSOR include dir: ${CUTENSOR_INCLUDE_DIR}")
 
   add_library(cutensor_binds SHARED)
 
@@ -335,6 +357,8 @@ if(BUILD_TESTING)
       tapp-reference-cutensor-demo
       PUBLIC
         ${CMAKE_CURRENT_SOURCE_DIR}/test
+      PRIVATE
+        ${CUTENSOR_INCLUDE_DIR}
     )
 
     add_test(
@@ -366,6 +390,12 @@ if(BUILD_TESTING)
       NAME tapp-reference-demo-dynamic
       COMMAND $<TARGET_FILE:tapp-reference-demo-dynamic>
     )
+
+    target_link_libraries(
+      tapp-reference-demo-dynamic
+      PRIVATE
+        ${CMAKE_DL_LIBS}
+    )
 
     # ----------------------------------------------------------------------------
     # test using dynamic library
@@ -390,6 +420,12 @@ if(BUILD_TESTING)
       NAME tapp-reference-test-dynamic
       COMMAND $<TARGET_FILE:tapp-reference-test-dynamic>
     )
+
+    target_link_libraries(
+      tapp-reference-test-dynamic
+      PRIVATE
+        ${CMAKE_DL_LIBS}
+    )
 
   endif()