Add CUDA sort shim for AOTI export (thrust-based sort_stable fallback)

Inductor emits aten::sort.stable for ops like argsort, but lacks a
native c-shim for it. This adds a thrust-based implementation
(aoti_torch_cuda_sort_stable) that handles int64, int32, and float32
dtypes on contiguous innermost-dim tensors. Registered as a supported
fallback kernel in CudaBackend so AOTI-compiled models can use sort.

This PR was authored with the assistance of Claude.

Author: digantdesai

backends/cuda/CMakeLists.txt

@@ -109,7 +109,9 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp
 
 # Only build int4mm shim when CUDA language/toolchain is available.
 if(CMAKE_CUDA_COMPILER)
-  list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu)
+  list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu
+       runtime/shims/sort.cu
+  )
 endif()
 
 add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources})

backends/cuda/cuda_backend.py

@@ -145,6 +145,7 @@ def save_data_externally(cls) -> bool:
     def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
         return {
             "at::_ops::_weight_int4pack_mm::call": None,
+            "at::_ops::sort_stable::call": None,
         }
 
     @classmethod

backends/cuda/runtime/TARGETS

@@ -33,13 +33,15 @@ runtime.cxx_library(
         "shims/cuda_guard.cpp",
         "shims/int4mm.cu",
         "shims/memory.cpp",
+        "shims/sort.cu",
         "shims/tensor_attribute.cpp",
     ],
     headers = [
         "shims/cuda_guard.h",
         "shims/int4mm.cuh",
         "shims/int4mm.h",
         "shims/memory.h",
+        "shims/sort.h",
         "shims/tensor_attribute.h",
         "utils.h",
     ],