feat(exla): configurable shm permissions for IPC buffers, default 0400 (#1732)

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 3 people

exla/Makefile

@@ -41,6 +41,10 @@ else
 	CFLAGS += -O3
 endif
 
+ifeq ($(MIX_ENV),prod)
+	CFLAGS += -DEXLA_PROD
+endif
+
 NVCC = $(CXX)
 NVCCFLAGS = $(CFLAGS)
 LDFLAGS += -L$(XLA_EXTENSION_LIB) -lxla_extension -shared -fvisibility=hidden

exla/c_src/exla/exla.cc

@@ -4,6 +4,7 @@
 #include <stdio.h>
 #include <string>
 #include <tuple>
+#include <unistd.h>
 #include <unordered_map>
 
 #include "absl/log/globals.h"
@@ -306,10 +307,10 @@ FINE_NIF(mlir_compile, ERL_NIF_DIRTY_JOB_CPU_BOUND);
 // ExlaBuffer Functions
 
 std::variant<std::tuple<fine::Atom, uint64_t, uint64_t>,
-             std::tuple<fine::Atom, std::string, uint64_t, uint64_t>,
              std::tuple<fine::Atom, std::string, uint64_t>>
 get_buffer_device_pointer(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
-                          fine::Term buffer_term, fine::Atom pointer_kind) {
+                          fine::Term buffer_term, fine::Atom pointer_kind,
+                          int64_t shm_permissions) {
   auto buffer = decode_exla_buffer(env, buffer_term);
 
   uint64_t device_size = unwrap(buffer->GetOnDeviceSizeInBytes());
@@ -321,22 +322,38 @@ get_buffer_device_pointer(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
 
   if (pointer_kind == "host_ipc") {
     auto handle_name =
-        "exla:ipc:" + std::to_string(device_size) + ":" + std::to_string(ptr);
-    auto fd = get_ipc_handle(handle_name.c_str(), device_size);
+        "exla:ipc:" + std::to_string(getpid()) + ":" + std::to_string(ptr);
+    auto fd = get_ipc_handle(handle_name.c_str(), device_size,
+                             static_cast<mode_t>(shm_permissions));
 
     if (fd == -1) {
       throw std::runtime_error("unable to get IPC handle");
     }
 
-    auto ipc_ptr = open_ipc_handle(fd, device_size);
+    auto ipc_ptr = open_ipc_handle(fd, device_size, /*writable=*/1);
     if (ipc_ptr == nullptr) {
       throw std::runtime_error("unable to open IPC handle");
     }
 
     memcpy(ipc_ptr, reinterpret_cast<void *>(ptr), device_size);
 
-    return std::make_tuple(pointer_kind, handle_name, static_cast<uint64_t>(fd),
-                           device_size);
+    // Repoint the original buffer at the shm mapping so both the exporter
+    // and any importers share the same physical pages.  This also frees
+    // the old XLA-managed memory, avoiding double memory usage.
+    auto shape = unwrap(buffer->buffer()->logical_on_device_shape());
+    auto device = unwrap(client->client()->LookupDevice(
+        xla::PjRtGlobalDeviceId(buffer->device_id())));
+    auto memory_space = unwrap(device->default_memory_space());
+
+    auto on_delete = [fd, ipc_ptr, device_size, handle_name]() {
+      close_ipc_handle(fd, ipc_ptr, handle_name.c_str(), device_size);
+    };
+
+    auto new_pjrt_buf = unwrap(client->client()->CreateViewOfDeviceBuffer(
+        ipc_ptr, shape, memory_space, on_delete));
+    buffer->ReplaceBuffer(std::move(new_pjrt_buf));
+
+    return std::make_tuple(pointer_kind, handle_name, device_size);
   }
 
   if (pointer_kind == "cuda_ipc") {
@@ -369,17 +386,20 @@ fine::ResourcePtr<ExlaBuffer> create_buffer_from_device_pointer(
     }
     ptr = maybe_pointer.value();
   } else if (pointer_kind == "host_ipc") {
-    auto tuple =
-        fine::decode<std::tuple<uint64_t, std::string>>(env, pointer_data);
-    auto fd = std::get<0>(tuple);
-    auto memname = std::get<1>(tuple);
+    auto memname = fine::decode<std::string>(env, pointer_data);
     auto device_size = xla::ShapeUtil::ByteSizeOf(shape);
-    ptr = open_ipc_handle(fd, device_size);
+    int writable = 0;
+    auto fd = open_existing_ipc_handle(memname.c_str(), &writable);
+    if (fd == -1) {
+      throw std::runtime_error("unable to get fd for IPC handle");
+    }
+    ptr = open_ipc_handle(fd, device_size, writable);
     if (ptr == nullptr) {
+      close(fd);
       throw std::runtime_error("unable to get pointer for IPC handle");
     }
-    on_delete_callback = [fd, memname, ptr, device_size]() {
-      close_ipc_handle(fd, ptr, memname.c_str(), device_size);
+    on_delete_callback = [fd, ptr, device_size]() {
+      close_imported_ipc_handle(fd, ptr, device_size);
     };
   } else if (pointer_kind == "local") {
     auto ptr_int = fine::decode<int64_t>(env, pointer_data);
@@ -677,6 +697,22 @@ fine::Ok<> start_log_sink(ErlNifEnv *env, ErlNifPid logger_pid) {
 
 FINE_NIF(start_log_sink, 0);
 
+// Test-only NIFs — excluded from production builds.
+#ifndef EXLA_PROD
+
+// Writes `data` into the memory at `address + offset`.  Intentionally unsafe:
+// no bounds checking.  The caller is responsible for ensuring the pointer is
+// valid and the region is large enough.
+fine::Ok<> write_to_pointer(ErlNifEnv *env, uint64_t address,
+                            ErlNifBinary data, uint64_t offset) {
+  uint8_t *ptr = reinterpret_cast<uint8_t *>(address);
+  std::memcpy(ptr + offset, data.data, data.size);
+  return fine::Ok();
+}
+FINE_NIF(write_to_pointer, 0);
+
+#endif // EXLA_PROD
+
 } // namespace exla
 
 FINE_INIT("Elixir.EXLA.NIF");

exla/c_src/exla/exla_client.cc

@@ -64,6 +64,20 @@ tsl::StatusOr<ERL_NIF_TERM> ExlaBuffer::ToBinary(ErlNifEnv *env,
   return binary_term;
 }
 
+void ExlaBuffer::ReplaceBuffer(std::unique_ptr<xla::PjRtBuffer> new_buffer) {
+  if (buffer_ && !buffer_->IsDeleted()) {
+    TrackDeallocation();
+    buffer_->Delete();
+  }
+  buffer_ = std::move(new_buffer);
+  if (client_ && buffer_) {
+    auto size_or = GetOnDeviceSizeInBytes();
+    if (size_or.ok()) {
+      client_->TrackBufferAllocated(device_id(), size_or.value());
+    }
+  }
+}
+
 tsl::Status ExlaBuffer::Deallocate() {
   if (buffer_->IsDeleted()) {
     return xla::FailedPrecondition(

exla/c_src/exla/exla_client.h

@@ -45,6 +45,10 @@ class ExlaBuffer {
 
   void SetClient(ExlaClient *client) { client_ = client; }
 
+  // Replace the underlying PjRt buffer with a new one (e.g. an shm-backed
+  // view).  The old buffer is deallocated first so XLA can free its memory.
+  void ReplaceBuffer(std::unique_ptr<xla::PjRtBuffer> new_buffer);
+
   ~ExlaBuffer();
 
 private:

exla/c_src/exla/ipc.cc

@@ -1,14 +1,17 @@
 #include "ipc.h"
 
 #include <cstdio>
+#include <errno.h>
 #include <fcntl.h>
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <unistd.h>
 
-// Function to create or open a shared memory object and set its size
-int get_ipc_handle(const char* memname, size_t memsize) {
-  int fd = shm_open(memname, O_CREAT | O_RDWR, 0666);
+// Create or open a shared memory object and set its size. `mode` is the
+// file mode bits forwarded to shm_open(3). The default is chosen in the
+// Elixir caller (EXLA.Backend.to_pointer/2).
+int get_ipc_handle(const char* memname, size_t memsize, mode_t mode) {
+  int fd = shm_open(memname, O_CREAT | O_RDWR, mode);
   if (fd == -1) {
     return -1;
   }
@@ -21,16 +24,36 @@ int get_ipc_handle(const char* memname, size_t memsize) {
   return fd;
 }
 
-// Function to map the shared memory in this process
-void* open_ipc_handle(int fd, size_t memsize) {
-  void* ptr = mmap(NULL, memsize, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+// Try O_RDWR first (permissions allow write); fall back to O_RDONLY on EACCES.
+int open_existing_ipc_handle(const char* memname, int* out_writable) {
+  int fd = shm_open(memname, O_RDWR, 0);
+  if (fd != -1) {
+    *out_writable = 1;
+    return fd;
+  }
+  if (errno == EACCES) {
+    fd = shm_open(memname, O_RDONLY, 0);
+    if (fd != -1) {
+      *out_writable = 0;
+      return fd;
+    }
+  }
+  return -1;
+}
+
+// MAP_SHARED zero-copy mapping. `writable` adds PROT_WRITE; fd access mode
+// must match (O_RDWR for writable, O_RDONLY for read-only).
+void* open_ipc_handle(int fd, size_t memsize, int writable) {
+  int prot = writable ? (PROT_READ | PROT_WRITE) : PROT_READ;
+  void* ptr = mmap(NULL, memsize, prot, MAP_SHARED, fd, 0);
   if (ptr == MAP_FAILED) {
     perror("mmap");
     return nullptr;
   }
   return ptr;
 }
 
+// Sender cleanup: remove the shm name so the object is freed once all fds close.
 int close_ipc_handle(int fd, void* ptr, const char* memname, size_t memsize) {
   if (munmap(ptr, memsize) == -1) {
     return -1;
@@ -44,3 +67,10 @@ int close_ipc_handle(int fd, void* ptr, const char* memname, size_t memsize) {
 
   return 0;
 }
+
+// Receiver cleanup: only munmap + close; the sender owns the shm name/lifetime.
+int close_imported_ipc_handle(int fd, void* ptr, size_t memsize) {
+  munmap(ptr, memsize);
+  close(fd);
+  return 0;
+}

exla/c_src/exla/ipc.h

@@ -1,7 +1,23 @@
 #pragma once
 
 #include <cstddef>
+#include <sys/types.h>
 
-int get_ipc_handle(const char* memname, size_t memsize);
-void* open_ipc_handle(int fd, size_t memsize);
+// Create a new shm segment, set its size, and return a writable fd.
+// `mode` is the file permission bits (e.g. 0o400, 0o600).
+int get_ipc_handle(const char* memname, size_t memsize, mode_t mode);
+
+// Open an existing shm segment.  Tries O_RDWR first; if EACCES, falls back to
+// O_RDONLY.  Sets *out_writable to 1 if write access was granted, 0 otherwise.
+// Returns -1 on error.
+int open_existing_ipc_handle(const char* memname, int* out_writable);
+
+// Map a shm fd with MAP_SHARED.  `writable` controls whether PROT_WRITE is
+// included; the fd must have been opened with the matching access mode.
+void* open_ipc_handle(int fd, size_t memsize, int writable);
+
+// Sender cleanup: munmap + close + shm_unlink.
 int close_ipc_handle(int fd, void* ptr, const char* memname, size_t memsize);
+
+// Receiver cleanup: munmap + close only (no unlink — sender owns the name).
+int close_imported_ipc_handle(int fd, void* ptr, size_t memsize);

exla/lib/exla/backend.ex

@@ -84,7 +84,14 @@ defmodule EXLA.Backend do
 
   @impl true
   def to_pointer(%T{data: %B{buffer: buffer}}, opts \\ []) do
-    opts = Keyword.validate!(opts, mode: :local)
+    opts = Keyword.validate!(opts, mode: :local, permissions: 0o400)
+    permissions = opts[:permissions]
+
+    unless is_integer(permissions) and permissions >= 0 and permissions <= 0o7777 do
+      raise ArgumentError,
+            ":permissions must be an integer in the range 0..0o7777 " <>
+              "(typically an octal literal like 0o400), got: #{inspect(permissions)}"
+    end
 
     mode =
       case {opts[:mode], buffer} do
@@ -108,18 +115,18 @@ defmodule EXLA.Backend do
 
         {mode, _} ->
           raise ArgumentError,
-                "expected one of :local, :cuda_ipc, :host_ipc, got: #{inspect(mode)}"
+                "expected one of :local, :ipc, got: #{inspect(mode)}"
       end
 
     client = EXLA.Client.fetch!(buffer.client_name)
 
-    case EXLA.NIF.get_buffer_device_pointer(client.ref, buffer.ref, mode) do
+    case EXLA.NIF.get_buffer_device_pointer(client.ref, buffer.ref, mode, permissions) do
       {:local, ptr, size} ->
         # Pointer is an integer here
         %Nx.Pointer{kind: :local, address: ptr, data_size: size}
 
-      {:host_ipc, handle_name, fd, size} ->
-        %Nx.Pointer{kind: :ipc, handle: handle_name, address: fd, data_size: size}
+      {:host_ipc, handle_name, size} ->
+        %Nx.Pointer{kind: :ipc, handle: handle_name, data_size: size}
 
       {:cuda_ipc, handle, size} ->
         %Nx.Pointer{kind: :ipc, handle: handle, address: buffer.device_id, data_size: size}
@@ -153,8 +160,8 @@ defmodule EXLA.Backend do
         {%Nx.Pointer{kind: :local, address: address}, _} ->
           {:local, address}
 
-        {%Nx.Pointer{kind: :ipc, address: fd, handle: handle}, :host} ->
-          {:host_ipc, {fd, handle}}
+        {%Nx.Pointer{kind: :ipc, handle: handle}, :host} ->
+          {:host_ipc, handle}
 
         {%Nx.Pointer{kind: :ipc, handle: handle}, :cuda} ->
           {:cuda_ipc, handle}

exla/lib/exla/nif.ex

@@ -55,7 +55,8 @@ defmodule EXLA.NIF do
   def mlir_get_typespec(_tensor), do: err!()
   def mlir_module_to_string(_builder), do: err!()
 
-  def get_buffer_device_pointer(_client, _buffer, _pointer_kind), do: err!()
+  def get_buffer_device_pointer(_client, _buffer, _pointer_kind, _shm_permissions),
+    do: err!()
 
   def create_buffer_from_device_pointer(
         _client,
@@ -96,5 +97,11 @@ defmodule EXLA.NIF do
   def encode_local_pid(_pid), do: err!()
   def decode_local_pid(_pid_bin), do: err!()
 
+  if Mix.env() != :prod do
+    # Writes `data` into the memory at `address + offset`.  Test-only; not
+    # compiled in production builds.
+    def write_to_pointer(_address, _data, _offset), do: err!()
+  end
+
   defp err!(), do: :erlang.nif_error(:undef)
 end