Import torch to wait for CUDA work in bindings built without CUDA support

Bindings built without CUDA support have no CUDA runtime of their own, so
they wait for CUDA work through torch. Before torch was imported lazily,
importing ExecuTorch imported it whenever it was installed. Now a wait
found torch missing unless something else had imported it, and failed
with an error that blamed the build. That broke CUDA inputs from other
libraries, such as CuPy, in 3 places: set_inputs() followed by
execute(), a call to a method that returns no tensor, and get_attribute()
of a CUDA attribute.

A wait now imports torch, unless EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT is
executorch, which keeps torch out. When torch is still not there, the
error says that imported torch would also work, and the note on a torch
that fails to import names the wait as a reason too.

Author: shoumikhin

extension/pybindings/README.md

@@ -157,9 +157,9 @@ device_input = ExecuTorchResult.from_host(
 Copying CUDA memory requires CUDA-enabled bindings: `from_host()`, `to_host()`
 of a CUDA result, and a CUDA `ExecuTorchResult` with `clone_outputs=True`.
 Bindings built without CUDA support synchronize through Python Torch instead,
-so once torch is imported, by you or by a call that returns torch tensors, they
-still read CUDA inputs, torch tensors or DLPack, and return borrowed CUDA
-results.
+and import torch to do so unless `EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT` is
+`executorch`. So with torch installed, or imported by you, they still read
+CUDA inputs, torch tensors or DLPack, and return borrowed CUDA results.
 
 Execution is synchronous at the Python boundary: a call returns after its device
 work has finished, so its outputs are ready on any stream. How long a call with

extension/pybindings/pybindings.cpp

@@ -805,8 +805,9 @@ bool torch_available() {
       if (!error.matches(PyExc_ImportError)) {
         try {
           error.value().attr("add_note")(
-              "ExecuTorch imports torch here to return torch.Tensor outputs. "
-              "Set EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch before "
+              "ExecuTorch imports torch here to return torch.Tensor outputs "
+              "or to wait for CUDA work. Set "
+              "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch before "
               "importing ExecuTorch to skip it.");
         } catch (...) {
           // add_note is best effort; never hide the original error.
@@ -845,7 +846,8 @@ bool keep_torch_outputs() {
 // then, keep_torch_outputs() keeps torch.Tensor where that rule would change
 // it. A call with no tensor inputs returns torch.Tensor whenever torch can be
 // imported, now and later, so its type never depends on what imported torch
-// first. A call that returns no tensor never imports torch.
+// first. Picking the type for a call that returns no tensor never imports
+// torch.
 PortableTensorOutput portable_tensor_output(
     bool returns_tensor,
     bool torch_input,
@@ -2369,6 +2371,18 @@ std::shared_ptr<DeviceMemoryBuffer> allocate_device_output(
   return std::make_shared<DeviceMemoryBuffer>(std::move(buffer.get()));
 }
 
+// Bindings built without CUDA support wait for CUDA work through torch, which
+// used to be imported with ExecuTorch. So a wait imports it, unless the
+// setting keeps torch out.
+void synchronize_device_for_call(int index) {
+#ifndef EXECUTORCH_PYBINDINGS_CUDA
+  if (tensor_output_override() != PortableTensorOutput::ExecuTorch) {
+    torch_available();
+  }
+#endif
+  synchronize_device(index);
+}
+
 // The Python boundary is synchronous, and each wait is for the whole device,
 // so one wait per device covers every tensor of a run. Before the run, another
 // library may still be writing a CUDA input on its own streams: a torch tensor,
@@ -2397,7 +2411,7 @@ void synchronize_tensor_devices(const Method& method) {
     add(method.get_output(i));
   }
   for (const auto index : devices) {
-    synchronize_device(index);
+    synchronize_device_for_call(index);
   }
 }
 
@@ -4879,7 +4893,7 @@ struct PyMethod final {
 #else
     // No execute() in this call has waited for the work that last wrote it.
     if (!attr->device().is_cpu() && attr->nbytes() != 0) {
-      synchronize_device(attr->device().index());
+      synchronize_device_for_call(attr->device().index());
     }
     // Mutable state lives in the method's planned memory. Anything else is a
     // constant, which can sit in read-only pages, so it is never returned as

extension/pybindings/pybindings_executorch_result.h

@@ -135,12 +135,15 @@ inline void synchronize_torch_cuda_device(int index) {
 // once torch is imported.
 inline void synchronize_device(int index) {
 #ifndef EXECUTORCH_PYBINDINGS_CUDA
-  if (loaded_torch_module() != nullptr) {
-    synchronize_torch_cuda_device(index);
-    return;
+  if (loaded_torch_module() == nullptr) {
+    throw std::runtime_error(
+        "CUDA tensor interchange requires bindings built with CUDA support, "
+        "or torch imported to wait for CUDA work");
   }
-#endif
+  synchronize_torch_cuda_device(index);
+#else
   synchronize_cuda_device(index);
+#endif
 }
 
 /** Tensor result returned by portable Python bindings. */

extension/pybindings/test/test_pybindings.py

@@ -4161,6 +4161,105 @@ def forward(self, x, y):
                 method.execute(cuda_inputs)[0].cpu(), torch.full((4, 4), 4.0)
             )
 
+    def test_cuda_input_imports_torch_to_wait_without_cuda_support(self):
+        if self.kernel_mode != "portable":
+            self.skipTest("only portable bindings wait for CUDA work themselves")
+        if not torch.cuda.is_available():
+            self.skipTest("needs a visible CUDA device")
+        from executorch.extension.pybindings import portable_lib
+
+        if "CudaBackend" in portable_lib._get_registered_backend_names():
+            self.skipTest("bindings built with CUDA support wait without torch")
+
+        from executorch.exir._serialize._program import PTEFile, serialize_pte_binary
+        from executorch.exir.schema import ExtraTensorInfo, Tensor
+
+        class Identity(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        program = to_edge(export(Identity(), (torch.ones(4),))).to_executorch(
+            ExecutorchBackendConfig(
+                memory_planning_pass=MemoryPlanningPass(
+                    alloc_graph_input=False, alloc_graph_output=False
+                )
+            )
+        )
+        schema = program.executorch_program
+        for entry in schema.execution_plan[0].values:
+            if isinstance(entry.val, Tensor):
+                entry.val.extra_tensor_info = ExtraTensorInfo(
+                    device_type=DeviceType.CUDA, device_index=0
+                )
+        data = bytes(serialize_pte_binary(PTEFile(schema)))
+        # Device memory from the CUDA driver, so nothing imports torch before
+        # the bindings wait for the device.
+        script = """
+import ctypes
+import sys
+
+from executorch.extension.pybindings import portable_lib
+
+class Device(ctypes.Structure):
+    _fields_ = [("device_type", ctypes.c_int32), ("device_id", ctypes.c_int32)]
+
+class DataType(ctypes.Structure):
+    _fields_ = [("code", ctypes.c_uint8), ("bits", ctypes.c_uint8), ("lanes", ctypes.c_uint16)]
+
+class DLTensor(ctypes.Structure):
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("device", Device),
+        ("ndim", ctypes.c_int32),
+        ("dtype", DataType),
+        ("shape", ctypes.POINTER(ctypes.c_int64)),
+        ("strides", ctypes.POINTER(ctypes.c_int64)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+class ManagedTensor(ctypes.Structure):
+    _fields_ = [("dl_tensor", DLTensor), ("manager_ctx", ctypes.c_void_p), ("deleter", ctypes.c_void_p)]
+
+driver = ctypes.CDLL("libcuda.so.1")
+context, pointer = ctypes.c_void_p(), ctypes.c_uint64()
+assert driver.cuInit(0) == 0
+assert driver.cuDevicePrimaryCtxRetain(ctypes.byref(context), 0) == 0
+assert driver.cuCtxSetCurrent(context) == 0
+assert driver.cuMemAlloc_v2(ctypes.byref(pointer), 16) == 0
+shape = (ctypes.c_int64 * 1)(4)
+managed = ManagedTensor(
+    DLTensor(pointer.value, Device(2, 0), 1, DataType(2, 32, 1), shape, None, 0)
+)
+make_capsule = ctypes.pythonapi.PyCapsule_New
+make_capsule.restype = ctypes.py_object
+make_capsule.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+
+with open(sys.argv[1], "rb") as source:
+    method = portable_lib._load_program_from_buffer(source.read()).load_method("forward")
+assert "torch" not in sys.modules
+method.set_inputs([make_capsule(ctypes.addressof(managed), b"dltensor", None)])
+method.execute()
+assert "torch" in sys.modules
+"""
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "cuda_identity.pte")
+            with open(path, "wb") as handle:
+                handle.write(data)
+            environment = os.environ.copy()
+            environment.pop("EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT", None)
+            completed = subprocess.run(
+                [sys.executable, "-c", script, path],
+                env=environment,
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}",
+        )
+
     def test_cuda_torch_tensor_interchange_without_aten(self):
         if self.kernel_mode != "portable":
             self.skipTest("only portable bindings use the Python Torch adapter")