CUDA backend: optional cuBLAS for matrix multiplications

backends/cuda/CMakeLists.txt

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -254,10 +254,17 @@
 )
 
 # Only build CUDA-only operator shims when the CUDA toolchain is available.
+set(_aoti_cuda_shims_use_cublas OFF)
 if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER)
   list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu
        runtime/shims/sort.cu runtime/shims/rand.cu
   )
+  # The cuBLAS matrix multiplication shims load cuBLASLt with dlopen and are
+  # not ported to Windows yet, so aoti_cuda_shims.lib must not list them.
+  if(NOT WIN32)
+    list(APPEND _aoti_cuda_shim_sources runtime/shims/mm.cu)
+    set(_aoti_cuda_shims_use_cublas ON)
+  endif()
 endif()
 
 add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources})
@@ -322,6 +329,15 @@
   )
 endif()
 
+# cuBLASLt is loaded on the first matrix multiplication rather than linked, so
+# models that never call the cuBLAS shims do not need the library installed.
+if(_aoti_cuda_shims_use_cublas)
+  target_compile_definitions(
+    aoti_cuda_shims
+    PRIVATE ET_CUBLASLT_LIBRARY="libcublasLt.so.${CUDAToolkit_VERSION_MAJOR}"
+  )
+endif()
+
 if(NOT _cuda_is_msvc_toolchain)
   executorch_target_link_options_shared_lib(aoti_cuda_shims)
 endif()

backends/cuda/cuda_backend.py

@@ -459,6 +459,29 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool:
     return value == "ON"
 
 
+def _gemm_backends(spec: Optional[CompileSpec], platform: str) -> str:
+    """Matrix multiplication choices autotuning may pick from: TRITON compiles
+    them into the model, ATEN calls cuBLAS at run time.
+
+    ATEN is opt-in because a cuBLAS call is a separate kernel that Inductor
+    cannot fuse with the ops around it.
+    """
+    if spec is None:
+        return "TRITON"
+    value = spec.value.decode("utf-8").upper().replace(" ", "")
+    backends = set(value.split(","))
+    if not backends <= {"ATEN", "TRITON"}:
+        raise ValueError(
+            f"Invalid {spec.key}: {value}. Expected ATEN, TRITON or ATEN,TRITON."
+        )
+    if "ATEN" in backends and (torch.version.hip is not None or platform == "windows"):
+        raise ValueError(
+            f"ATEN in {spec.key} runs matrix multiplications on cuBLAS, which "
+            "only the Linux CUDA runtime provides."
+        )
+    return value
+
+
 CUDA_GRAPH_AUTOTUNE_TIMING_COMPILE_SPEC = "cuda_graph_autotune_timing"
 
 # The CUDA tensors `move_program_to_device` placed on the GPU for the compile on
@@ -773,6 +796,9 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
             "at::_ops::sort_stable::call": None,
             "aoti_torch_cuda_sort_stable": None,
             "aoti_torch_cuda_randint_low_out": None,
+            "aoti_torch_cuda_mm_out": None,
+            "aoti_torch_cuda_addmm_out": None,
+            "aoti_torch_cuda_bmm_out": None,
         }
 
     @classmethod
@@ -857,8 +883,6 @@ def get_aoti_compile_options(
             ),
             # Enable maximum automatic tuning for optimal performance
             "max_autotune": True,
-            # Use TRITON for GEMM (General Matrix Multiply) operations tuning only to avoid using operators in libtorch
-            "max_autotune_gemm_backends": "TRITON",
             # Use TRITON backend for convolution operations tuning only to avoid using operators in libtorch
             "max_autotune_conv_backends": "TRITON",
             "aot_inductor.emit_multi_arch_kernel": emit_multi_arch_kernel,
@@ -870,6 +894,7 @@ def get_aoti_compile_options(
         emulate_precision_casts = True
         max_autotune = True
         autotune_at_compile_time = None
+        gemm_backends_spec = None
         shim_library_path = None
         for spec in compile_specs:
             if spec.key == "platform":
@@ -882,6 +907,11 @@ def get_aoti_compile_options(
                 autotune_at_compile_time = _on_off_compile_spec_value(spec)
             elif spec.key == "shim_library_path":
                 shim_library_path = spec.value.decode("utf-8")
+            elif spec.key == "max_autotune_gemm_backends":
+                gemm_backends_spec = spec
+        options["max_autotune_gemm_backends"] = _gemm_backends(
+            gemm_backends_spec, platform
+        )
         options["emulate_precision_casts"] = emulate_precision_casts
         options["max_autotune"] = max_autotune
         if autotune_at_compile_time is not None:

backends/cuda/runtime/aoti_cuda_shims.lib.md

@@ -33,4 +33,5 @@ the same source differ.
 
 A name only resolves if something in the DLL defines it. The DLL is built from the
 CUDA shims plus the SlimTensor common shims, so a shim added to the ETensor common
-shims will link and then fail to load.
+shims will link and then fail to load. The cuBLAS shims in `shims/mm.cu` are not
+built for Windows, so they must not be listed either.