Cover caller-stream TLS semantics

backends/aoti/CMakeLists.txt

@@ -117,6 +117,13 @@ target_compile_options(
 target_compile_definitions(
   aoti_common_shims_slim PUBLIC $<$<PLATFORM_ID:Windows>:EXPORT_AOTI_FUNCTIONS>
 )
+# Private executorch_ prefix on the slim aoti_torch_* shims so they never
+# collide with libtorch's in a coalesced process (see backends/aoti/export.h).
+# Must match aot_inductor.shim_symbol_prefix set in
+# backends/cuda/cuda_backend.py.
+target_compile_definitions(
+  aoti_common_shims_slim PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_
+)
 
 target_link_libraries(aoti_common_shims_slim PUBLIC slimtensor ${CMAKE_DL_LIBS})
 
@@ -139,6 +146,9 @@ if(MSVC)
     aoti_common_shims_slim_obj
     PUBLIC $<$<PLATFORM_ID:Windows>:EXPORT_AOTI_FUNCTIONS>
   )
+  target_compile_definitions(
+    aoti_common_shims_slim_obj PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_
+  )
   target_link_libraries(
     aoti_common_shims_slim_obj PUBLIC slimtensor ${CMAKE_DL_LIBS}
   )

backends/aoti/aoti_delegate_handle.h

@@ -10,6 +10,7 @@
 
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/freeable_buffer.h>
 #include <string>
 
@@ -19,7 +20,11 @@ namespace aoti {
 
 using executorch::runtime::Error;
 using executorch::runtime::FreeableBuffer;
-using executorch::runtime::etensor::Tensor;
+// Mode-flexible host-tensor spelling: `etensor::Tensor` in portable mode (the
+// slim ETensor the AOTI blob passes around), `at::Tensor` under USE_ATEN_LIB.
+// Used here only as an opaque handle type; the CUDA delegate reinterpret_casts
+// its SlimTensor* through it, and Metal passes its own portable handle.
+using executorch::aten::Tensor;
 
 extern "C" {
 

backends/aoti/export.h

@@ -24,3 +24,104 @@
 #else
 #define AOTI_SHIM_EXPORT
 #endif
+
+// Keep ExecuTorch's aoti_torch_* shim *definitions* in lockstep with the symbol
+// names an AOTInductor blob imports. The blob and these (libtorch-free) shim
+// libraries are both compiled with -DAOTI_SHIM_SYMBOL_PREFIX=executorch_ (see
+// cuda_backend.py get_aoti_compile_options and the shim targets'
+// exported_preprocessor_flags), so this block renames the exported definitions
+// to executorch_aoti_torch_*. The blob (whose imports torch's
+// aoti_torch/c/macros.h renames identically) then binds to ExecuTorch's
+// SlimTensor shims and never to libtorch's at::Tensor aoti_torch_* in a
+// coalesced (TensorRT[ATen] + CUDA) process.
+//
+// This block is kept self-contained (no torch include) to preserve the slim
+// layer's libtorch-free property. It MUST stay in sync with the identical list
+// in torch's aoti_torch/c/shim_symbol_prefix.h; a mismatch is a loud link error
+// (an unresolved executorch_aoti_torch_* symbol), never a silent crash. With
+// the define unset (portable builds) it expands to nothing.
+#ifdef AOTI_SHIM_SYMBOL_PREFIX
+#define AOTI_SHIM_CONCAT2_(a, b) a##b
+#define AOTI_SHIM_CONCAT_(a, b) AOTI_SHIM_CONCAT2_(a, b)
+#define AOTI_SHIM_RENAME_(name) AOTI_SHIM_CONCAT_(AOTI_SHIM_SYMBOL_PREFIX, name)
+
+#define aoti_torch_assign_tensors_out \
+  AOTI_SHIM_RENAME_(aoti_torch_assign_tensors_out)
+#define aoti_torch_check AOTI_SHIM_RENAME_(aoti_torch_check)
+#define aoti_torch_clone AOTI_SHIM_RENAME_(aoti_torch_clone)
+#define aoti_torch_clone_preserve_strides \
+  AOTI_SHIM_RENAME_(aoti_torch_clone_preserve_strides)
+#define aoti_torch_copy_ AOTI_SHIM_RENAME_(aoti_torch_copy_)
+#define aoti_torch_create_cuda_guard \
+  AOTI_SHIM_RENAME_(aoti_torch_create_cuda_guard)
+#define aoti_torch_create_cuda_stream_guard \
+  AOTI_SHIM_RENAME_(aoti_torch_create_cuda_stream_guard)
+#define aoti_torch_create_tensor_from_blob \
+  AOTI_SHIM_RENAME_(aoti_torch_create_tensor_from_blob)
+#define aoti_torch_create_tensor_from_blob_v2 \
+  AOTI_SHIM_RENAME_(aoti_torch_create_tensor_from_blob_v2)
+#define aoti_torch_cuda_guard_set_index \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_guard_set_index)
+#define aoti_torch_cuda_int4_plain_mm \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_int4_plain_mm)
+#define aoti_torch_cuda_int5_plain_mm \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_int5_plain_mm)
+#define aoti_torch_cuda_int6_plain_mm \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_int6_plain_mm)
+#define aoti_torch_cuda_int8_plain_mm \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_int8_plain_mm)
+#define aoti_torch_cuda_rand AOTI_SHIM_RENAME_(aoti_torch_cuda_rand)
+#define aoti_torch_cuda_randint_low_out \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_randint_low_out)
+#define aoti_torch_cuda_sort_stable \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda_sort_stable)
+#define aoti_torch_cuda__weight_int4pack_mm \
+  AOTI_SHIM_RENAME_(aoti_torch_cuda__weight_int4pack_mm)
+#define aoti_torch_delete_cuda_guard \
+  AOTI_SHIM_RENAME_(aoti_torch_delete_cuda_guard)
+#define aoti_torch_delete_cuda_stream_guard \
+  AOTI_SHIM_RENAME_(aoti_torch_delete_cuda_stream_guard)
+#define aoti_torch_delete_tensor_object \
+  AOTI_SHIM_RENAME_(aoti_torch_delete_tensor_object)
+#define aoti_torch_device_type_cpu AOTI_SHIM_RENAME_(aoti_torch_device_type_cpu)
+#define aoti_torch_device_type_cuda \
+  AOTI_SHIM_RENAME_(aoti_torch_device_type_cuda)
+#define aoti_torch_dtype_bfloat16 AOTI_SHIM_RENAME_(aoti_torch_dtype_bfloat16)
+#define aoti_torch_dtype_bool AOTI_SHIM_RENAME_(aoti_torch_dtype_bool)
+#define aoti_torch_dtype_float16 AOTI_SHIM_RENAME_(aoti_torch_dtype_float16)
+#define aoti_torch_dtype_float32 AOTI_SHIM_RENAME_(aoti_torch_dtype_float32)
+#define aoti_torch_dtype_int16 AOTI_SHIM_RENAME_(aoti_torch_dtype_int16)
+#define aoti_torch_dtype_int32 AOTI_SHIM_RENAME_(aoti_torch_dtype_int32)
+#define aoti_torch_dtype_int64 AOTI_SHIM_RENAME_(aoti_torch_dtype_int64)
+#define aoti_torch_dtype_int8 AOTI_SHIM_RENAME_(aoti_torch_dtype_int8)
+#define aoti_torch_dtype_uint8 AOTI_SHIM_RENAME_(aoti_torch_dtype_uint8)
+#define aoti_torch_empty_strided AOTI_SHIM_RENAME_(aoti_torch_empty_strided)
+#define aoti_torch_empty_strided_pinned \
+  AOTI_SHIM_RENAME_(aoti_torch_empty_strided_pinned)
+#define aoti_torch_get_current_cuda_stream \
+  AOTI_SHIM_RENAME_(aoti_torch_get_current_cuda_stream)
+#define aoti_torch_get_data_ptr AOTI_SHIM_RENAME_(aoti_torch_get_data_ptr)
+#define aoti_torch_get_device_index \
+  AOTI_SHIM_RENAME_(aoti_torch_get_device_index)
+#define aoti_torch_get_device_type AOTI_SHIM_RENAME_(aoti_torch_get_device_type)
+#define aoti_torch_get_dim AOTI_SHIM_RENAME_(aoti_torch_get_dim)
+#define aoti_torch_get_dtype AOTI_SHIM_RENAME_(aoti_torch_get_dtype)
+#define aoti_torch_get_numel AOTI_SHIM_RENAME_(aoti_torch_get_numel)
+#define aoti_torch_get_sizes AOTI_SHIM_RENAME_(aoti_torch_get_sizes)
+#define aoti_torch_get_storage_offset \
+  AOTI_SHIM_RENAME_(aoti_torch_get_storage_offset)
+#define aoti_torch_get_storage_size \
+  AOTI_SHIM_RENAME_(aoti_torch_get_storage_size)
+#define aoti_torch_get_strides AOTI_SHIM_RENAME_(aoti_torch_get_strides)
+#define aoti_torch_grad_mode_is_enabled \
+  AOTI_SHIM_RENAME_(aoti_torch_grad_mode_is_enabled)
+#define aoti_torch_grad_mode_set_enabled \
+  AOTI_SHIM_RENAME_(aoti_torch_grad_mode_set_enabled)
+#define aoti_torch_item_bool AOTI_SHIM_RENAME_(aoti_torch_item_bool)
+#define aoti_torch_layout_strided AOTI_SHIM_RENAME_(aoti_torch_layout_strided)
+#define aoti_torch_new_tensor_handle \
+  AOTI_SHIM_RENAME_(aoti_torch_new_tensor_handle)
+#define aoti_torch__reinterpret_tensor \
+  AOTI_SHIM_RENAME_(aoti_torch__reinterpret_tensor)
+#define aoti_torch_warn AOTI_SHIM_RENAME_(aoti_torch_warn)
+#endif // AOTI_SHIM_SYMBOL_PREFIX

backends/aoti/shim_symbol_prefix.bzl

@@ -0,0 +1,21 @@
+# Single source of truth for the ExecuTorch AOTI shim symbol prefix.
+#
+# The AOTInductor blob generated by the CUDA delegate is compiled with
+# -DAOTI_SHIM_SYMBOL_PREFIX=<prefix> (set from Python via
+# cuda_backend.get_aoti_compile_options -> aot_inductor.shim_symbol_prefix), and
+# ExecuTorch's shim libraries must be compiled with the SAME define so their
+# exported aoti_torch_* definitions rename to <prefix>aoti_torch_* and satisfy
+# the blob's now-prefixed imports. Keeping the value here (and mirroring it in
+# the Python default) prevents blob-import / shim-export drift, which would
+# surface as a loud unresolved-symbol link error rather than a runtime crash.
+#
+# See:
+#   - fbcode/executorch/backends/aoti/export.h (shim-export rename block)
+#   - fbcode/caffe2/torch/csrc/inductor/aoti_torch/c/shim_symbol_prefix.h
+#   - fbcode/executorch/backends/cuda/cuda_backend.py (Python side)
+
+AOTI_SHIM_SYMBOL_PREFIX = "executorch_"
+
+AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS = [
+    "-DAOTI_SHIM_SYMBOL_PREFIX=" + AOTI_SHIM_SYMBOL_PREFIX,
+]

backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp

@@ -12,6 +12,7 @@
 #include <executorch/runtime/platform/platform.h>
 #include <gtest/gtest.h>
 
+#include <thread>
 #include <type_traits>
 
 using namespace executorch::backends::cuda;
@@ -296,6 +297,30 @@ TEST(CallerStreamGuardTest, GuardSelectsThenRestores) {
   EXPECT_FALSE(getCallerStream().has_value());
 }
 
+TEST(CallerStreamGuardTest, ExplicitNullStreamIsStillSelected) {
+  CallerStreamGuard guard(nullptr);
+  ASSERT_TRUE(getCallerStream().has_value());
+  EXPECT_EQ(*getCallerStream(), nullptr);
+}
+
+TEST(CallerStreamGuardTest, SelectionIsThreadLocal) {
+  const cudaStream_t selected = fake_stream(0);
+  CallerStreamGuard guard(selected);
+
+  bool child_started_without_stream = false;
+  bool child_selected_own_stream = false;
+  std::thread child([&]() {
+    child_started_without_stream = !getCallerStream().has_value();
+    CallerStreamGuard child_guard(fake_stream(1));
+    child_selected_own_stream = getCallerStream() == fake_stream(1);
+  });
+  child.join();
+
+  EXPECT_TRUE(child_started_without_stream);
+  EXPECT_TRUE(child_selected_own_stream);
+  EXPECT_EQ(getCallerStream(), selected);
+}
+
 TEST(CallerStreamGuardTest, NestedGuardsRestoreOuter) {
   const cudaStream_t outer = fake_stream(1);
   const cudaStream_t inner = fake_stream(2);

backends/aoti/targets.bzl

@@ -1,5 +1,6 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load("@fbsource//tools/build_defs:fbsource_utils.bzl", "is_fbcode")
+load(":shim_symbol_prefix.bzl", "AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS")
 
 def define_common_targets():
     if not is_fbcode():
@@ -105,6 +106,18 @@ def define_common_targets():
             "export.h",
             "utils.h",
         ],
+        # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole)
+        # These extern "C" aoti_torch_* shims are only referenced by the CUDA
+        # delegate's dlopen'd AOTInductor blob, which is invisible to the static
+        # linker; without link_whole, --gc-sections drops the whole TU and the
+        # blob fails to resolve them at load.
+        link_whole = True,
+        # Export aoti_torch_* under the executorch_ prefix (see export.h) so the
+        # blob binds to these SlimTensor shims, not libtorch's, in a coalesced
+        # process. In BOTH lists: preprocessor_flags renames this target's own
+        # definitions (common_shims_slim.cpp); exported_ renames header includers.
+        preprocessor_flags = AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS,
+        exported_preprocessor_flags = AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS,
         visibility = ["@EXECUTORCH_CLIENTS"],
         exported_deps = [
             "//executorch/runtime/core:core",

backends/cuda/CMakeLists.txt

@@ -126,6 +126,15 @@ add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources})
 
 # Define CUDA_AVAILABLE to use SlimTensor on GPU in common_shims_slim.h
 target_compile_definitions(aoti_cuda_shims PRIVATE CUDA_AVAILABLE=1)
+# Private executorch_ prefix on the slim aoti_torch_* shims so they never
+# collide with libtorch's in a coalesced process (see backends/aoti/export.h).
+# Must match aot_inductor.shim_symbol_prefix set in
+# backends/cuda/cuda_backend.py. PUBLIC so direct consumers that include the
+# shim headers (e.g. the shim gtests) compile the declarations under the same
+# prefix and link the DLL's prefixed exports.
+target_compile_definitions(
+  aoti_cuda_shims PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_
+)
 
 # Define export macros for shared library. Use WIN32 (not just MSVC) so MinGW
 # cross-compiles also emit dllexport symbols for AOTI shims.

backends/cuda/cuda_backend.py

@@ -431,6 +431,25 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
             passes.append(ReplaceEdgeOpWithTritonOpPass())
         return passes
 
+    @staticmethod
+    def _apply_shim_symbol_prefix(options: Dict[str, typing.Any]) -> None:
+        """Give the generated blob a private aoti_torch_* ABI so it binds to
+        ExecuTorch's SlimTensor shims, never libtorch's at::Tensor shims, in a
+        coalesced process (e.g. a TensorRT[ATen] + CUDA .pte). Must match the
+        -DAOTI_SHIM_SYMBOL_PREFIX the ExecuTorch shim libs are built with
+        (backends/aoti/shim_symbol_prefix.bzl). The shim libs export only the
+        prefixed names, so this PyTorch capability is required, not optional; a
+        missing key means the installed torch predates it.
+        """
+        if not hasattr(torch._inductor.config.aot_inductor, "shim_symbol_prefix"):
+            raise RuntimeError(
+                "The installed PyTorch does not support "
+                "aot_inductor.shim_symbol_prefix, which the ExecuTorch CUDA "
+                "backend requires to match its prefixed aoti_torch_* shims. "
+                "Update PyTorch to a version that provides it."
+            )
+        options["aot_inductor.shim_symbol_prefix"] = "executorch_"
+
     @classmethod
     def get_aoti_compile_options(
         cls, compile_specs: List[CompileSpec]
@@ -469,29 +488,29 @@ def get_aoti_compile_options(
             "aot_inductor.emit_multi_arch_kernel": emit_multi_arch_kernel,
         }
 
-        try:
-            import torch
+        cls._apply_shim_symbol_prefix(options)
 
+        try:
             options["aot_inductor.custom_ops_to_c_shims"] = {
                 torch.ops.executorch_cuda.int4_plain_mm.default: [
-                    "AOTITorchError aoti_torch_cuda_int4_plain_mm("
+                    "AOTITorchError executorch_aoti_torch_cuda_int4_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "int64_t, AtenTensorHandle*)"
                 ],
                 torch.ops.executorch_cuda.int5_plain_mm.default: [
-                    "AOTITorchError aoti_torch_cuda_int5_plain_mm("
+                    "AOTITorchError executorch_aoti_torch_cuda_int5_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, int64_t, AtenTensorHandle*)"
                 ],
                 torch.ops.executorch_cuda.int6_plain_mm.default: [
-                    "AOTITorchError aoti_torch_cuda_int6_plain_mm("
+                    "AOTITorchError executorch_aoti_torch_cuda_int6_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, AtenTensorHandle, int64_t, AtenTensorHandle*)"
                 ],
                 torch.ops.executorch_cuda.int8_plain_mm.default: [
-                    "AOTITorchError aoti_torch_cuda_int8_plain_mm("
+                    "AOTITorchError executorch_aoti_torch_cuda_int8_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, int64_t, AtenTensorHandle*)"
                 ],