diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index efdfaab00cd..8dc20ee5eed 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -471,6 +471,21 @@ def _tensor_spec_to_evalue( ctypes.c_char * typing.cast(torch.UntypedStorage, spec.storage).nbytes() ) + # data_ptr() is dereferenced from the host below; for non-CPU storage + # (e.g. CUDA) that read can segfault instead of raising (#21545). + storage_device = typing.cast(torch.UntypedStorage, spec.storage).device + if spec.allocated_memory != 0 and storage_device.type != "cpu": + raise InternalError( + self._emit_node_specific_error( + self.node, + f"Constant tensor storage is on device '{storage_device}', " + "but constants serialized into the program must be in CPU " + "memory. Move them to CPU before export (e.g. model.cpu()), " + "or delegate the ops that use them to a backend that " + "accepts device tensors.", + ) + ) + buffer_data = ( bytes( ctypes.cast( diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index 78480d3b55c..c52dbf32276 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -1833,6 +1833,15 @@ def forward(self, x): self.assertEqual(len(program.constant_buffer[1].storage), 4) self.assertEqual(len(program.constant_buffer[2].storage), 4) + def test_emit_non_cpu_constant_raises(self) -> None: + # Like CUDA storage, "meta" storage is not host-readable, but needs no + # accelerator. Unguarded, it already fails with a wrapped "NULL pointer + # access" InternalError, so match on the device name. + model = torch.nn.Linear(4, 3).to("meta") + edge = to_edge(export(model, (torch.ones(2, 4, device="meta"),), strict=True)) + with self.assertRaisesRegex(InternalError, "on device 'meta'"): + edge.to_executorch() + def test_non_persistent_buffer(self) -> None: class NonPersistentBuffer(nn.Module): def __init__(self):