Add cuSolverDx JIT fusion and solver projections (#1176)

* Add cuSolverDx JIT fusion and solver projections

Add support for cuSolverDx within JIT path. Solver functions are different than most operators since many of them return multiple values, and this was not compatible with the MatX syntax in the past. We introduce a new projection syntax where multiple values can be used inside of a single statement. This works because in non-JIT mode the `mtie` syntax is used and there's no ambiguity, but in JIT mode none of the overhead from the operator is present since it's converted to a string.

Author: cliffburdick

AGENTS.md

@@ -9,3 +9,13 @@ the @.devcontainer directory using the create_base_container.sh script. The pyth
 ## Compiling and Running
 For use in external projects, MatX simply needs to be included by `#include <matx.h>`. To build unit tests the use the CMake option `MATX_BUILD_TESTS=ON`, benchmarks 
 `MATX_BUILD_BENCHMARKS=ON`, and examples `MATX_BUILD_EXAMPLES=ON`. Individual tests can be compiled separately via different targets output from CMake.
+
+## Development Expectations
+Every new public function, operator, transform, or backend path should include accompanying unit tests and documentation updates. Tests should cover the supported
+types, ranks, batching behavior, and important error or fallback cases for the new behavior. Documentation should describe how to use the feature, required build
+options or dependencies, and any limitations that would affect users or future maintainers. Any change to operators should also verify
+@docs_input/executor_compatibility.rst and update it if executor support, limitations, or backend requirements changed.
+
+When adding or changing accelerated backends such as CUDA, cuBLAS, cuSolver, cuFFT, or MathDx paths, preserve the existing non-accelerated behavior unless the task
+explicitly calls for a breaking change. Prefer focused tests that compare the new backend against an existing trusted MatX path, and include negative tests for
+unsupported shape, dtype, rank, or configuration combinations when applicable.

CMakeLists.txt

@@ -159,6 +159,9 @@ rapids_cpm_cccl(
 )
 
 target_link_libraries(matx INTERFACE CCCL::CCCL)
+if(CCCL_SOURCE_DIR)
+    target_compile_definitions(matx INTERFACE MATX_CCCL_SOURCE_DIR="${CCCL_SOURCE_DIR}")
+endif()
 
 # Set flags for compiling tests faster (only for nvcc)
 if (NOT CMAKE_CUDA_COMPILER_ID STREQUAL "Clang")
@@ -364,13 +367,28 @@ if (MATX_EN_JIT OR MATX_EN_MATHDX)
 
     # Link NVRTC library
     target_link_libraries(matx INTERFACE CUDA::nvrtc)
+    if(TARGET CUDA::nvJitLink)
+        target_link_libraries(matx INTERFACE CUDA::nvJitLink)
+    endif()
 endif()
 
 if (MATX_EN_MATHDX)
-    set(MathDx_VERSION 25.06)
+    set(MathDx_VERSION 26.03)
     set(MathDx_NANO 0)
+    set(LIBMATHDX_VERSION 0.3.2)
     include(cmake/FindMathDx.cmake)
     target_compile_definitions(matx INTERFACE MATX_EN_MATHDX)
+    target_compile_definitions(matx INTERFACE
+        MATX_MATHDX_INCLUDE_DIR="${MATHDX_INCLUDE_DIR}"
+        MATX_MATHDX_CUTLASS_INCLUDE_DIR="${MATHDX_CUTLASS_INCLUDE_DIR}"
+        MATX_LIBMATHDX_INCLUDE_DIR="${LIBMATHDX_INCLUDE_DIR}"
+    )
+    if(EXISTS "${MATHDX_ROOT}/lib/libcusolverdx.fatbin")
+        target_compile_definitions(matx INTERFACE MATX_CUSOLVERDX_FATBIN="${MATHDX_ROOT}/lib/libcusolverdx.fatbin")
+    endif()
+    if(EXISTS "${MATHDX_ROOT}/lib/libcusolverdx.a")
+        target_compile_definitions(matx INTERFACE MATX_CUSOLVERDX_LIBRARY="${MATHDX_ROOT}/lib/libcusolverdx.a")
+    endif()
 
     # Link libmathdx if available
     if(TARGET libmathdx::libmathdx)
@@ -379,7 +397,7 @@ if (MATX_EN_MATHDX)
     endif()
 
     # Link mathdx components
-    target_link_libraries(matx INTERFACE mathdx::cufftdx)
+    target_link_libraries(matx INTERFACE mathdx::cufftdx mathdx::cublasdx_no_lto mathdx::cusolverdx)
 endif()
 
 if (MATX_EN_CUDSS)

cmake/FindMathDx.cmake

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2021, NVIDIA CORPORATION.
+# Copyright (c) 2021-2026, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,84 +14,25 @@
 # limitations under the License.
 #=============================================================================
 
-#[=======================================================================[.rst:
-FindMathDx
---------
-
-Find MathDx
-
-Imported targets
-^^^^^^^^^^^^^^^^
-
-This module defines the following :prop_tgt:`IMPORTED` target(s):
-
-``MathDx::MathDx``
-  The MathDx library, if found.
-
-Result variables
-^^^^^^^^^^^^^^^^
-
-This module will set the following variables in your project:
-
-``MathDx_FOUND``
-  True if MathDx is found.
-``MathDx_INCLUDE_DIRS``
-  The include directories needed to use MathDx.
-``MathDx_VERSION_STRING``
-  The version of the MathDx library found. [OPTIONAL]
-
-#]=======================================================================]
 set(MathDx_VERSION_FULL ${MathDx_VERSION}.${MathDx_NANO})
 
-# Prefer using a Config module if it exists for this project
-set(MathDx_NO_CONFIG FALSE)
-if(NOT MathDx_NO_CONFIG)
-  find_package(MathDx CONFIG QUIET HINTS ${MathDx_DIR})
-  if(MathDx_FOUND)
-    find_package_handle_standard_args(MathDx DEFAULT_MSG MathDx_CONFIG)
-    return()
-  endif()
+if(NOT CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
+  message(FATAL_ERROR "MATX_EN_MATHDX requires CUDA 13.0 or newer for MathDx ${MathDx_VERSION_FULL}")
 endif()
 
-find_path(MathDx_INCLUDE_DIR NAMES MathDx.h)
+set(MATHDX_CUDA_VERSION "cuda13")
+set(MATHDX_CUDA_SUFFIX "cuda13.0")
 
-# Search for the MathDx library
-find_library(MathDx_LIBRARY 
-  NAMES MathDx mathdx
-  HINTS ${MathDx_DIR}
-  PATH_SUFFIXES lib lib64
-)
-
-include(${CMAKE_ROOT}/Modules/FindPackageHandleStandardArgs.cmake)
-
-find_package_handle_standard_args(MathDx
-                                  REQUIRED_VARS MathDx_LIBRARY MathDx_INCLUDE_DIR
-                                  VERSION_VAR )
-
-if(NOT MathDx_FOUND)
-  set(MathDx_FILENAME libMathDx-linux-x86_64-${MathDx_VERSION}-archive)
-
-  message(STATUS "MathDx not found. Downloading library. By continuing this download you accept to the license terms of MathDx")
-
-  CPMAddPackage(
-    NAME MathDx
-    VERSION ${MathDx_VERSION}
-    URL https://developer.download.nvidia.com/compute/cuFFTDx/redist/cuFFTDx/nvidia-mathdx-${MathDx_VERSION_FULL}.tar.gz
-    DOWNLOAD_ONLY YES 
-  )
-endif()
+message(STATUS "Using MathDx ${MathDx_VERSION_FULL} (${MATHDX_CUDA_VERSION})")
+message(STATUS "Using libmathdx ${LIBMATHDX_VERSION} (${MATHDX_CUDA_VERSION})")
 
-# Download libmathdx based on CUDA version and platform
-# Detect CUDA version (12 or 13)
-if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
-  set(LIBMATHDX_CUDA_VERSION "cuda13")
-  set(LIBMATHDX_CUDA_SUFFIX "cuda13.0")
-else()
-  set(LIBMATHDX_CUDA_VERSION "cuda12")
-  set(LIBMATHDX_CUDA_SUFFIX "cuda12.0")
-endif()
+CPMAddPackage(
+  NAME MathDx
+  VERSION ${MathDx_VERSION_FULL}
+  URL https://developer.nvidia.com/downloads/compute/cuSOLVERDx/redist/cuSOLVERDx/${MATHDX_CUDA_VERSION}/nvidia-mathdx-${MathDx_VERSION_FULL}-${MATHDX_CUDA_VERSION}.tar.gz
+  DOWNLOAD_ONLY YES
+)
 
-# Detect platform
 if(WIN32)
   set(LIBMATHDX_PLATFORM "win32-x86_64")
   set(LIBMATHDX_EXT "zip")
@@ -103,63 +44,51 @@ elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
   endif()
   set(LIBMATHDX_EXT "tar.gz")
 else()
-  message(WARNING "Unsupported platform for libmathdx download")
+  message(FATAL_ERROR "Unsupported platform for libmathdx download")
 endif()
 
-# Set libmathdx version
-set(LIBMATHDX_VERSION "0.2.3")
-
-# Download libmathdx if platform is supported
-if(DEFINED LIBMATHDX_PLATFORM)
-  set(LIBMATHDX_URL "https://developer.nvidia.com/downloads/compute/cublasdx/redist/cublasdx/${LIBMATHDX_CUDA_VERSION}/libmathdx-${LIBMATHDX_PLATFORM}-${LIBMATHDX_VERSION}-${LIBMATHDX_CUDA_SUFFIX}.${LIBMATHDX_EXT}")
-  
-  message(STATUS "Downloading libmathdx for ${LIBMATHDX_PLATFORM} with ${LIBMATHDX_CUDA_VERSION}")
-  message(STATUS "libmathdx URL: ${LIBMATHDX_URL}")
-  
-  CPMAddPackage(
-    NAME libmathdx
-    VERSION ${LIBMATHDX_VERSION}
-    URL ${LIBMATHDX_URL}
-    DOWNLOAD_ONLY YES
-  )
-  
-  # Add libmathdx to the search paths
-  set(LIBMATHDX_ROOT "${PROJECT_BINARY_DIR}/_deps/libmathdx-src")
-  list(APPEND CMAKE_PREFIX_PATH "${LIBMATHDX_ROOT}")
-  
-  # Find libmathdx library file
-  find_library(LIBMATHDX_LIBRARY
-    NAMES mathdx libmathdx
-    PATHS "${LIBMATHDX_ROOT}/lib"
-    NO_DEFAULT_PATH
+set(LIBMATHDX_URL "https://developer.nvidia.com/downloads/compute/cublasdx/redist/cublasdx/${MATHDX_CUDA_VERSION}/libmathdx-${LIBMATHDX_PLATFORM}-${LIBMATHDX_VERSION}-${MATHDX_CUDA_SUFFIX}.${LIBMATHDX_EXT}")
+
+message(STATUS "libmathdx URL: ${LIBMATHDX_URL}")
+
+CPMAddPackage(
+  NAME libmathdx
+  VERSION ${LIBMATHDX_VERSION}
+  URL ${LIBMATHDX_URL}
+  DOWNLOAD_ONLY YES
+)
+
+set(MATX_MATHDX_ROOT "${PROJECT_BINARY_DIR}/_deps/mathdx-src/nvidia/mathdx/${MathDx_VERSION}")
+set(MATHDX_INCLUDE_DIR "${MATX_MATHDX_ROOT}/include")
+set(MATHDX_CUTLASS_INCLUDE_DIR "${MATX_MATHDX_ROOT}/external/cutlass/include")
+set(LIBMATHDX_ROOT "${PROJECT_BINARY_DIR}/_deps/libmathdx-src")
+set(LIBMATHDX_INCLUDE_DIR "${LIBMATHDX_ROOT}/include")
+
+find_library(LIBMATHDX_LIBRARY
+  NAMES mathdx libmathdx
+  PATHS "${LIBMATHDX_ROOT}/lib"
+  NO_DEFAULT_PATH
+)
+
+if(NOT LIBMATHDX_LIBRARY OR NOT EXISTS "${LIBMATHDX_INCLUDE_DIR}/libmathdx.h")
+  message(FATAL_ERROR "Could not find libmathdx ${LIBMATHDX_VERSION} library or headers after download")
+endif()
+
+if(NOT TARGET libmathdx::libmathdx)
+  add_library(libmathdx::libmathdx INTERFACE IMPORTED)
+  set_target_properties(libmathdx::libmathdx PROPERTIES
+    INTERFACE_INCLUDE_DIRECTORIES "${LIBMATHDX_INCLUDE_DIR}"
+    INTERFACE_LINK_LIBRARIES "${LIBMATHDX_LIBRARY}"
   )
-  
-  # Set include directories (in both local and parent scope)
-  set(LIBMATHDX_INCLUDE_DIR "${LIBMATHDX_ROOT}/include")
-  set(LIBMATHDX_INCLUDE_DIR "${LIBMATHDX_INCLUDE_DIR}" PARENT_SCOPE)
-  
-  if(LIBMATHDX_LIBRARY AND EXISTS ${LIBMATHDX_INCLUDE_DIR})
-    message(STATUS "Found libmathdx library: ${LIBMATHDX_LIBRARY}")
-    message(STATUS "Found libmathdx include dir: ${LIBMATHDX_INCLUDE_DIR}")
-    
-    # Create libmathdx target
-    if(NOT TARGET libmathdx::libmathdx)
-      add_library(libmathdx::libmathdx INTERFACE IMPORTED)
-      set_target_properties(libmathdx::libmathdx PROPERTIES
-        INTERFACE_INCLUDE_DIRECTORIES "${LIBMATHDX_INCLUDE_DIR}"
-        INTERFACE_LINK_LIBRARIES "${LIBMATHDX_LIBRARY}"
-      )
-    endif()
-  else()
-    message(WARNING "Could not find libmathdx library or include directory after download")
-  endif()
 endif()
 
-find_package(mathdx REQUIRED COMPONENTS cufftdx CONFIG
-PATHS
-    "${PROJECT_BINARY_DIR}/_deps/mathdx-src/nvidia/mathdx/${MathDx_VERSION}/lib/cmake/mathdx/"
-    "${PROJECT_BINARY_DIR}/_deps/libmathdx-src/lib/cmake/libmathdx/"
-    "${PROJECT_BINARY_DIR}/_deps/libmathdx-src"
-    "/opt/nvidia/mathdx/${MathDx_VERSION_FULL}"
+set(cublasdx_CUTLASS_ROOT "${MATX_MATHDX_ROOT}/external/cutlass")
+set(cusolverdx_CUTLASS_ROOT "${MATX_MATHDX_ROOT}/external/cutlass")
+
+find_package(mathdx REQUIRED COMPONENTS cufftdx cublasdx cusolverdx CONFIG
+  PATHS
+    "${MATX_MATHDX_ROOT}/lib/cmake/mathdx"
+    "/opt/nvidia/mathdx/${MathDx_VERSION}"
 )
 
+set(MATHDX_ROOT "${MATX_MATHDX_ROOT}")

docs_input/api/linalg/decomp/chol.rst

@@ -8,6 +8,12 @@ Perform a Cholesky factorization.
 .. note::
   The input matrix must be symmetric positive-definite
 
+.. note::
+  CUDA JIT kernel fusion is supported by cuSolverDx if ``-DMATX_EN_MATHDX=ON`` is enabled. The current JIT path
+  supports rank 2 through 4 square matrices with ``float``, ``double``, ``complex<float>``, and
+  ``complex<double>`` values. Unsupported ranks, shapes, or data types should use a normal executor and will be
+  rejected by ``CUDAJITExecutor``.
+
 .. versionadded:: 0.6.0
 
 .. doxygenfunction:: chol

docs_input/api/linalg/decomp/inverse.rst

@@ -11,6 +11,19 @@ Compute the inverse of a square matrix.
 .. note::
    This function is currently not supported with host-based executors (CPU)
 
+.. note::
+   CUDA JIT kernel fusion is supported by cuSolverDx if ``-DMATX_EN_MATHDX=ON`` is enabled. The current JIT path
+   supports rank 2 through 4 square matrices with ``float``, ``double``, ``complex<float>``, and
+   ``complex<double>`` values. Unsupported ranks, shapes, or data types should use a normal executor and will be
+   rejected by ``CUDAJITExecutor``.
+
+.. note::
+   The default inverse uses LU/GETRF-style factorization for general square matrices. For Hermitian
+   positive-definite inputs, ``inv<MAT_INVERSE_ALGO_POSV>(A)`` enables a cuSolverDx POSV JIT path that solves
+   ``A * X = I`` directly and can fuse with other compatible JIT operators such as cuBLASDx matmul. POSV is
+   currently a MathDx/CUDAJITExecutor-only inverse algorithm, and callers must ensure the input satisfies the
+   positive-definite contract.
+
 .. versionadded:: 0.6.0
 
 .. doxygenfunction:: inv(const OpA &a)
@@ -23,5 +36,3 @@ Examples
    :start-after: example-begin inv-test-1
    :end-before: example-end inv-test-1
    :dedent:
-
-

docs_input/api/linalg/decomp/lu.rst

@@ -7,6 +7,12 @@ Perform an LU factorization.
 
 .. versionadded:: 0.6.0
 
+.. note::
+   The ``mtie`` assignment form of ``lu`` uses the normal non-JIT solver path. CUDA JIT fusion is available through
+   lazy projection members such as ``lu(A).LU`` and ``lu(A).Piv`` when ``-DMATX_EN_MATHDX=ON`` is enabled and the
+   runtime shape/type is supported by cuSolverDx. Projection JIT currently supports ranks 2 through 4 and
+   ``float``, ``double``, ``complex<float>``, and ``complex<double>`` inputs.
+
 .. doxygenfunction:: lu
 
 Examples
@@ -17,3 +23,15 @@ Examples
    :start-after: example-begin lu-test-1
    :end-before: example-end lu-test-1
    :dedent:
+
+Projection Examples
+~~~~~~~~~~~~~~~~~~~
+
+Lazy projections let ``LU`` and ``Piv`` participate in a larger expression. This example is included from the
+``CuSolverDxPivotProjectionUsedInFusedExpression`` unit test.
+
+.. literalinclude:: ../../../../test/00_solver/LU.cu
+   :language: cpp
+   :start-after: example-begin lu-projection-test-1
+   :end-before: example-end lu-projection-test-1
+   :dedent: