emlx

defmodule EMLX.MixProject do
  use Mix.Project

  @app :emlx
  @version "0.4.2"
  @mlx_version "0.32.0"
  @source_url "https://github.com/elixir-nx/emlx"
  # mlx-build Darwin archives from v0.31.0+ encode the deployment target in the
  # filename. 26.2 ships AOT NAX (Metal 4); 14.0 is the portable reduced set.
  @macos_deployment_target_latest "26.2"
  @macos_deployment_target_compat "14.0"
  @macos_deployment_targets [@macos_deployment_target_compat, @macos_deployment_target_latest]

  require Logger

  def project do
    libmlx_config = libmlx_config()

    [
      app: @app,
      version: @version,
      description: "MLX bindings and backend for Nx",
      elixir: "~> 1.17",
      elixirc_paths: elixirc_paths(Mix.env()),
      start_permanent: Mix.env() == :prod,
      deps: deps(),
      docs: docs(),
      # elixir_make — a function ref so Fine.include_dir/0 (needs the :fine
      # dep compiled) isn't called while this project/0 map is being built.
      make_env: fn ->
        %{
          "MLX_DIR" => libmlx_config.dir,
          "MLX_VERSION" => libmlx_config.version,
          "MLX_BUILD" => to_string(libmlx_config.features.build?),
          "MLX_INCLUDE_DIR" =>
            Path.join(
              libmlx_config.dir,
              if(libmlx_config.features.build?, do: "usr/include", else: "include")
            ),
          "MLX_LIB_DIR" =>
            Path.join(
              libmlx_config.dir,
              if(libmlx_config.features.build?, do: "usr/lib", else: "lib")
            ),
          "MLX_VARIANT" => libmlx_config.variant,
          "EMLX_CACHE_DIR" => libmlx_config.cache_dir,
          "EMLX_VERSION" => @version,
          "LIBMLX_ENABLE_DEBUG" => to_string(libmlx_config.features.debug?),
          "FINE_INCLUDE_DIR" => Fine.include_dir(),
          "MACOSX_DEPLOYMENT_TARGET" => libmlx_config.macos_deployment_target
        }
      end,

      # Compilers
      compilers: compilers(libmlx_config),
      aliases: aliases(),
      package: package()
    ]
  end

  def application do
    [
      extra_applications: [:logger, :inets, :ssl, :public_key, :crypto],
      mod: {EMLX.Application, []}
    ]
  end

  def cli do
    [preferred_envs: [docs: :docs, "hex.publish": :docs]]
  end

  defp elixirc_paths(:test), do: ["lib", "test/support"]
  defp elixirc_paths(_), do: ["lib"]

  defp deps do
    [
      {:elixir_make, "~> 0.6"},
      {:fine, "~> 0.1", runtime: false},
      {:nx, "~> 0.13.1"},
      {:telemetry, "~> 1.0"},
      {:ex_doc, "~> 0.34", only: :docs}
    ]
  end

  defp docs do
    [
      main: "EMLX",
      source_url_pattern: "#{@source_url}/blob/v#{@version}/emlx/%{path}#L%{line}",
      extras: [
        "README.md"
      ]
    ]
  end

  defp aliases do
    [
      "compile.mlx": &download_and_unarchive/1
    ]
  end

  defp compilers(libmlx_config) do
    compilers = [:elixir_make] ++ Mix.compilers()

    unless libmlx_config.features.build? do
      [:mlx] ++ compilers
    else
      compilers
    end
  end

  defp current_target_from_env do
    arch = System.get_env("TARGET_ARCH")
    os = System.get_env("TARGET_OS")
    abi = System.get_env("TARGET_ABI")

    if !Enum.all?([arch, os, abi], &Kernel.is_nil/1) do
      "#{arch}-#{os}-#{abi}"
    end
  end

  defp current_target! do
    case current_target() do
      {:ok, target} ->
        target

      {:error, reason} ->
        Mix.raise(reason)
    end
  end

  defp current_target do
    current_target_from_env = current_target_from_env()

    if current_target_from_env do
      # overwrite current target triplet from environment variables
      {:ok, current_target_from_env}
    else
      current_target(:os.type())
    end
  end

  defp current_target({:win32, _}) do
    processor_architecture =
      String.downcase(String.trim(System.get_env("PROCESSOR_ARCHITECTURE")))

    # https://docs.microsoft.com/en-gb/windows/win32/winprog64/wow64-implementation-details?redirectedfrom=MSDN
    partial_triplet =
      case processor_architecture do
        "amd64" ->
          "x86_64-windows-"

        "ia64" ->
          "ia64-windows-"

        "arm64" ->
          "aarch64-windows-"

        "x86" ->
          "x86-windows-"
      end

    {compiler, _} = :erlang.system_info(:c_compiler_used)

    case compiler do
      :msc ->
        {:ok, partial_triplet <> "msvc"}

      :gnuc ->
        {:ok, partial_triplet <> "gnu"}

      other ->
        {:ok, partial_triplet <> Atom.to_string(other)}
    end
  end

  defp current_target({:unix, _}) do
    # get current target triplet from `:erlang.system_info/1`
    system_architecture = to_string(:erlang.system_info(:system_architecture))
    current = String.split(system_architecture, "-", trim: true)

    case length(current) do
      4 ->
        {:ok, "#{Enum.at(current, 0)}-#{Enum.at(current, 2)}-#{Enum.at(current, 3)}"}

      3 ->
        case :os.type() do
          {:unix, :darwin} ->
            current =
              if "aarch64" == Enum.at(current, 0) do
                ["arm64" | tl(current)]
              else
                current
              end

            # could be something like aarch64-apple-darwin21.0.0
            # but we don't really need the last 21.0.0 part
            if String.match?(Enum.at(current, 2), ~r/^darwin.*/) do
              {:ok, "#{Enum.at(current, 0)}-#{Enum.at(current, 1)}-darwin"}
            else
              {:ok, system_architecture}
            end

          _ ->
            {:ok, system_architecture}
        end

      _ ->
        {:error, "Cannot determine current target"}
    end
  end

  @supported_targets [
    "x86_64-apple-darwin",
    "arm64-apple-darwin",
    "x86_64-linux-gnu",
    "aarch64-linux-gnu",
    "riscv64-linux-gnu"
  ]
  defp libmlx_config do
    version = System.get_env("LIBMLX_VERSION", @mlx_version)

    features = %{
      jit?: to_boolean(System.get_env("LIBMLX_ENABLE_JIT")),
      debug?: to_boolean(System.get_env("LIBMLX_ENABLE_DEBUG")),
      build?: to_boolean(System.get_env("LIBMLX_BUILD"))
    }

    current_target = current_target!()

    macos_deployment_target =
      unless features.build? do
        resolve_macos_deployment_target(current_target)
      end

    # Deployment target is a Darwin archive-name segment (mlx-build), not a
    # sorted feature flag — keep it ahead of `-debug`/`-jit`.
    variant = to_variant(features, macos_deployment_target)

    cache_dir =
      if dir = System.get_env("LIBMLX_CACHE") do
        Path.expand(dir)
      else
        :filename.basedir(:user_cache, "libmlx")
      end

    libmlx_archive =
      Path.join(
        cache_dir,
        "libmlx-#{version}-#{current_target}#{variant}.tar.gz"
      )

    libmlx_archive = System.get_env("MLX_ARCHIVE_PATH", libmlx_archive)

    features =
      if not Enum.member?(@supported_targets, current_target) and
           is_nil(System.get_env("MLX_ARCHIVE_PATH")) do
        Logger.warning("""
        Current target #{current_target} is not officially supported by EMLX, will fallback to building from source.

        A prebuilt libmlx archive for this target can be specified by setting the environment variable MLX_ARCHIVE_PATH to the path of the archive.
        """)

        %{features | build?: true}
      else
        features
      end

    %{
      target: current_target,
      libmlx_archive: libmlx_archive,
      version: version,
      dir: Path.join(cache_dir, "libmlx-#{version}-#{current_target}#{variant}"),
      features: features,
      variant: variant,
      macos_deployment_target: macos_deployment_target,
      cache_dir: cache_dir
    }
  end

  defp to_boolean(nil), do: false

  defp to_boolean(var) do
    String.downcase(to_string(var)) in ["1", "true", "on", "yes", "y"]
  end

  defp to_variant(features, macos_deployment_target) do
    feature_variant =
      [
        if(features.build?, do: "build"),
        if(features.debug?, do: "debug"),
        if(features.jit?, do: "jit")
      ]
      |> Enum.filter(&(&1 != nil))
      |> Enum.sort()
      |> Enum.join("-")

    feature_variant = if feature_variant != "", do: "-#{feature_variant}", else: ""

    case macos_deployment_target do
      nil -> feature_variant
      deployment_target -> "-#{deployment_target}#{feature_variant}"
    end
  end

  defp resolve_macos_deployment_target(target) do
    if String.contains?(target, "apple-darwin") do
      selected = selected_macos_deployment_target()
      assert_macos_deployment_target_compatible!(selected)
      selected
    end
  end

  defp selected_macos_deployment_target do
    case System.get_env("LIBMLX_DEPLOYMENT_TARGET") do
      nil ->
        if to_boolean(System.get_env("LIBMLX_MACOS_COMPAT")) do
          @macos_deployment_target_compat
        else
          @macos_deployment_target_latest
        end

      target when target in @macos_deployment_targets ->
        target

      other ->
        Mix.raise("""
        Invalid LIBMLX_DEPLOYMENT_TARGET=#{inspect(other)}.

        Supported values: #{Enum.join(@macos_deployment_targets, ", ")}.

        Use #{@macos_deployment_target_latest} (default) for the full macOS feature set \
        including AOT NAX kernels, or #{@macos_deployment_target_compat} / \
        LIBMLX_MACOS_COMPAT=true for the reduced-featureset build that runs on macOS 14+.
        """)
    end
  end

  defp assert_macos_deployment_target_compatible!(deployment_target) do
    # Cross-compiling via TARGET_* — host sw_vers is not the runtime OS.
    if current_target_from_env() do
      :ok
    else
      case macos_product_version() do
        nil ->
          :ok

        host_version ->
          min_version = parse_dotted_version!(deployment_target)

          if version_lt?(host_version, min_version) do
            Mix.raise("""
            Precompiled libmlx for macOS deployment target #{deployment_target} cannot run \
            on this host (macOS #{format_dotted_version(host_version)}).

            This is the failure mode behind https://github.com/elixir-nx/emlx/issues/127 \
            (VM crash: mutex lock failed).

            For macOS 15 and other hosts older than #{@macos_deployment_target_latest}, use \
            the reduced-featureset archive:

                export LIBMLX_MACOS_COMPAT=true

            or set LIBMLX_DEPLOYMENT_TARGET=#{@macos_deployment_target_compat} explicitly.
            """)
          else
            :ok
          end
      end
    end
  end

  defp macos_product_version do
    case :os.type() do
      {:unix, :darwin} ->
        case System.cmd("sw_vers", ["-productVersion"], stderr_to_stdout: true) do
          {version, 0} -> parse_dotted_version!(String.trim(version))
          _ -> nil
        end

      _ ->
        nil
    end
  end

  defp parse_dotted_version!(version) do
    case String.split(version, ".", trim: true) do
      [major, minor | _rest] ->
        {String.to_integer(major), String.to_integer(minor)}

      [major] ->
        {String.to_integer(major), 0}
    end
  rescue
    ArgumentError ->
      Mix.raise("Could not parse macOS version #{inspect(version)}")
  end

  defp version_lt?({maj_a, min_a}, {maj_b, min_b}) do
    maj_a < maj_b or (maj_a == maj_b and min_a < min_b)
  end

  defp format_dotted_version({major, minor}), do: "#{major}.#{minor}"

  defp download_and_unarchive(args) do
    libmlx_config = libmlx_config()

    if "--force" in args do
      File.rm_rf(libmlx_config.dir)
      File.rm_rf(libmlx_config.cache_dir)
    end

    if File.dir?(libmlx_config.dir) do
      {:ok, []}
    else
      download_and_unarchive(libmlx_config.cache_dir, libmlx_config)
    end
  end

  defp download_and_unarchive(cache_dir, libmlx_config) do
    File.mkdir_p!(cache_dir)

    libmlx_archive = libmlx_config.libmlx_archive

    url =
      "https://github.com/cocoa-xu/mlx-build/releases/download/v#{libmlx_config.version}/mlx-#{libmlx_config.target}#{libmlx_config.variant}.tar.gz"

    sha256_url = "#{url}.sha256"

    verify_integrity = "sha256=url:#{sha256_url}"

    {url, verify_integrity} =
      if customized_url = System.get_env("MLX_ARCHIVE_URL") do
        verify_integrity = System.get_env("MLX_ARCHIVE_INTEGRITY")
        {customized_url, verify_integrity}
      else
        {url, verify_integrity}
      end

    unless File.exists?(libmlx_archive) do
      # Download libmlx

      case :os.type() do
        {:unix, :darwin} ->
          :ok

        {:unix, _} ->
          Logger.warning("MLX only has CPU backend available for current target")

        _ ->
          Mix.raise("EMLX only supports macOS and x86_64, aarch64 and riscv64 Linux for now")
      end

      Mix.shell().info("Downloading libmlx from #{url}")
      download!(url, libmlx_archive)
      :ok = maybe_verify_integrity!(verify_integrity, libmlx_archive)
    end

    # Unpack libmlx and move to the target cache dir
    parent_libmlx_dir = Path.join(Path.dirname(libmlx_config.dir), "libmlx")
    File.mkdir_p!(parent_libmlx_dir)

    # Extract to the parent directory (it will be inside the libmlx directory)
    :ok = :erl_tar.extract(libmlx_archive, [:compressed, {:cwd, parent_libmlx_dir}])

    # And then rename
    File.rename!(parent_libmlx_dir, libmlx_config.dir)

    :ok
  end

  defp maybe_verify_integrity!(nil, _libmlx_archive), do: :ok

  defp maybe_verify_integrity!(verify_integrity, libmlx_archive) do
    {checksum_algo, expected_checksum} = get_checksum_info!(verify_integrity)
    libmlx_archive_checksum = checksum!(libmlx_archive, checksum_algo)

    if expected_checksum != libmlx_archive_checksum do
      Mix.raise(
        "Checksum (#{checksum_algo}) mismatch for #{libmlx_archive}. Expected: #{expected_checksum}, got: #{libmlx_archive_checksum}"
      )
    else
      :ok
    end
  end

  @known_checksum_algos [
    "sha",
    "sha224",
    "sha256",
    "sha384",
    "sha512",
    "sha3_224",
    "sha3_256",
    "sha3_384",
    "sha3_512",
    "blake2b",
    "blake2s",
    "ripemd160",
    "md4",
    "md5"
  ]

  defp get_checksum_info!(verify_integrity) do
    case String.split(verify_integrity, "=", parts: 2, trim: true) do
      [algo, checksum] when algo in @known_checksum_algos ->
        {String.to_existing_atom(algo), get_checksum_value!(checksum)}

      _ ->
        Mix.raise("Invalid checksum: #{verify_integrity}")
    end
  end

  defp get_checksum_value!("url:" <> url) do
    checksum_from_url!(url)
  end

  defp get_checksum_value!(checksum) do
    checksum
  end

  defp checksum_from_url!(url) do
    data = download!(url)
    checksum = String.split(data, " ", parts: 2, trim: true)

    if length(checksum) == 0 do
      Mix.raise("Invalid checksum file: #{url}")
    end

    hd(checksum)
  end

  def download!(url, save_as \\ nil) do
    url_charlist = String.to_charlist(url)

    if proxy = System.get_env("HTTP_PROXY") || System.get_env("http_proxy") do
      Mix.shell().info("Using HTTP_PROXY: #{proxy}")
      %{host: host, port: port} = URI.parse(proxy)

      :httpc.set_options([{:proxy, {{String.to_charlist(host), port}, []}}])
    end

    if proxy = System.get_env("HTTPS_PROXY") || System.get_env("https_proxy") do
      Mix.shell().info("Using HTTPS_PROXY: #{proxy}")
      %{host: host, port: port} = URI.parse(proxy)
      :httpc.set_options([{:https_proxy, {{String.to_charlist(host), port}, []}}])
    end

    # https://erlef.github.io/security-wg/secure_coding_and_deployment_hardening/inets
    # TODO: This may no longer be necessary from Erlang/OTP 25.0 or later.
    https_options = [
      ssl:
        [
          verify: :verify_peer,
          customize_hostname_check: [
            match_fun: :public_key.pkix_verify_hostname_match_fun(:https)
          ]
        ] ++ cacerts_options()
    ]

    options = [body_format: :binary]

    case :httpc.request(:get, {url_charlist, []}, https_options, options) do
      {:ok, {{_, 200, _}, _headers, body}} ->
        if save_as do
          File.write!(save_as, body)
        end

        body

      other ->
        Mix.raise("Failed to download #{url}, reason: #{inspect(other)}")
    end
  end

  defp cacerts_options do
    cond do
      path = System.get_env("HEX_CACERTS_PATH") ->
        [cacertfile: path]

      certs = otp_cacerts() ->
        [cacerts: certs]

      true ->
        warn_no_cacerts()
        []
    end
  end

  defp otp_cacerts do
    if System.otp_release() >= "25" do
      # cacerts_get/0 raises if no certs found
      try do
        :public_key.cacerts_get()
      rescue
        _ -> nil
      end
    end
  end

  defp warn_no_cacerts do
    Mix.shell().error("""
    No certificate trust store was found.

    A certificate trust store is required in
    order to download locales for your configuration.
    Since elixir_make could not detect a system
    installed certificate trust store one of the
    following actions may be taken:

    1. Specify the location of a certificate trust store
       by configuring it in environment variable:

         export HEX_CACERTS_PATH="/path/to/cacerts.pem"

    2. Use OTP 25+ on an OS that has built-in certificate
       trust store.
    """)
  end

  defp checksum!(file_path, algo) do
    case File.read(file_path) do
      {:ok, content} ->
        Base.encode16(:crypto.hash(algo, content), case: :lower)

      {:error, reason} ->
        Mix.raise(
          "Cannot read the file for checksum comparison: #{file_path}. Reason: #{inspect(reason)}"
        )
    end
  end

  defp package do
    [
      licenses: ["Apache-2.0"],
      links: %{"GitHub" => @source_url},
      maintainers: ["Paulo Valente", "Cocoa Xu", "Samrat Man Singh"]
    ]
  end
end