Packages

exla

0.7.2
0.13.0 0.12.0 0.11.0 0.10.0 0.9.2 0.9.1 0.9.0 0.8.0 0.7.3 0.7.2 0.7.1 0.7.0 0.6.4 0.6.3 0.6.1 0.6.0 0.5.3 0.5.2 0.5.1 0.5.0 0.4.2 0.4.1 0.4.0 0.3.0 0.2.3 0.2.2 0.2.1 0.2.0

Google's XLA (Accelerated Linear Algebra) compiler/backend for Nx

HexDocs

Current section

Files

Jump to
Readme
28 Versions
5 Dependencies
Dependants
Files
Activity
exla lib exla mlir context_pool.ex
Raw

lib/exla/mlir/context_pool.ex 

defmodule EXLA.MLIR.ContextPool do

  @moduledoc false

  # Internal pool for MLIRContext reference management

  @behaviour NimblePool



  def checkout(fun) when is_function(fun, 1) do

    NimblePool.checkout!(

      __MODULE__,

      :checkout,

      fn _pool, context -> {fun.(context), :ok} end,

      :infinity

    )

  end



  @impl NimblePool

  def init_worker(pool_state) do

    {:ok, context} = EXLA.NIF.new_mlir_context()

    {:ok, context, pool_state}

  end



  @impl NimblePool

  def handle_checkout(:checkout, _from, context, pool_state) do

    {:ok, context, context, pool_state}

  end



  @impl NimblePool

  def handle_checkin(:ok, _from, context, pool_state) do

    # We just keep the references around and let them die out upon worker termination/GC

    {:ok, context, pool_state}

  end



  @impl NimblePool

  def terminate_worker(_reason, _context, pool_state) do

    # GC will clean it up

    {:ok, pool_state}

  end

end