SME日記その18 SMEが使えるかどうかをElixirから判定する

Elixirにおける，SMEが使えるかどうかの判定ロジックを考えてみました．

SMEシリーズ

• Apple Silicon M4はM3シリーズからScalable Matrix Extension (SME)命令などが足されている
• SME日記その1: Apple Silicon M4に搭載されたScalable Matrix Extension(SME)のベクトル長(SVL)を取得する
• SME日記その2: Apple Silicon M4にはCVTW命令は備わっていない？
• SME日記その3: Apple Silicon M4にどの命令が実装されているかをsysctl hwの実行結果とドキュメントから推測する
• SME日記その4 Streaming SVE modeでCNTWを実行してみる．
• SME日記その5 Streaming SVE modeでCNTWを実行してみる Part 2
• SME日記その6 Streaming SVE modeでsvcntw()とsvcntsw()を実行してみる
• SME日記その7 svcntw()とRDSVL命令の実行結果の関係性を考察する
• SME日記その8 __arm_new("za")について調べる
• SME日記その9 OpenBLASのSME対応状況について調べる
• SME日記その10 Streaming SVE modeでCNTWを実行してみる(再考)
• SME日記その11 OpenBLASのSSCALでSMEが使われているかを検証してみる Part.1
• SME日記その12 OpenBLASのSSCALでSMEが使われているかを検証してみる Part.2
• SME日記その13 OpenBLASのSSCALでSMEが使われているかを検証してみる Part.3
• SME日記その14 AppleBLASのSSCALでSMEが使われているかを検証してみる Part.1
• SME日記その15 AppleBLASのSGEMMでSMEが使われているかを検証してみる Part.1
• SME日記その16 Scalable Matrix Extension (SME)の研究の今後の展望についての技術的ポエム
• SME日記その17 __arm_new("za")について調べる Part.2

ElixirでのSMEが使えるかどうかの判定ロジック

1. macOSでなければfalseを返して終了
2. sysctl hw.optional.arm を実行し，出力にhw.optional.arm.FEAT_SME: 1とhw.optional.arm.FEAT_SME2: 1が含まれていれば，2の結果はtrue，そうでなければ2の結果はfalse
3. 次のCプログラムを/usr/bin/clang -O2 -march=armv9-a+sme test_sme.c -o test_smeとしてコンパイルし，成功すれば3の結果はtrue，そうでなければ3の結果はfalse
4. 2と3の結果が共にtrueであれば，SMEを使えるのでtrueを返す．そうでなければfalseを返す．

1・2と，3 は並列に実行しても良い．

#include <arm_sme.h>

__arm_locally_streaming
__arm_new("za")
void test_arm_new(void)
{
}

int main(int argc, char *argv[])
{
  test_arm_new();
}

macOSかどうかの判定

case :os.type() do
  {:unix, :darwin} -> true
  _ -> false
end

sysctl hw.optional.arm の実行判定

case :os.type() do
  {:unix, :darwin} ->
    case System.find_executable("sysctl") do
      nil -> false

      executable -> 
        case System.cmd(executable, ["hw.optional.arm"], into: []) do
          {result, 0} ->
            result 
            |> Enum.map(&String.trim/1)
            |> Enum.filter(&String.match?(&1, ~r/FEAT\_SME/))
            |> Enum.map(&String.split(&1, " "))
            |> Enum.filter(fn [_, v] -> v == "1" end)
            |> Enum.count()
            |> then(& &1 != 0)

          {_, _} -> false
        end
    end
  
  _ -> false
end

Cプログラムのコンパイル

case System.find_executable("/usr/bin/clang") do
  nil -> false

  executable -> 
    base_root = :crypto.strong_rand_bytes(10) |> Base.encode32(case: :lower)
    base = base_root <> ".c"

    File.write(
      Path.join("/tmp", base), 
      """
      #include <arm_sme.h>

      __arm_locally_streaming
      __arm_new("za")
      void test_arm_new(void) {}
      int main(int argc, char *argv[])
      {
        test_arm_new();
      }
      """
    )
    |> case do
      :ok -> 
        System.cmd(
          executable, 
          [
            "-O2", 
            "-Werror",
            "-Wall",
            "-march=armv9-a+sme", 
            base,
            "-o",
            base_root
          ],
          into: [],
          cd: "/tmp",
          stderr_to_stdout: true
        )
        |> case do
          {_, 0} -> true
          _ -> false
        end
      
      {:error, _reason} -> false
    end
end

全体

defmodule SME do
  def available?() do
    {:ok, pid} = Task.Supervisor.start_link()

    task1 = Task.Supervisor.async(pid, fn -> runnable?() end)
    task2 = Task.Supervisor.async(pid, fn -> compilable?() end)
    
    Task.await(task1) and Task.await(task2)
  end

  def runnable?() do
    case :os.type() do
      {:unix, :darwin} -> 
        case execute("sysctl", ["hw.optional.arm"]) do
          {result, 0} -> 
            result
            |> String.split("\n")
            |> Enum.map(&String.trim/1)
            |> Enum.filter(&String.match?(&1, ~r/FEAT\_SME/))
            |> Enum.map(&String.split(&1, " "))
            |> Enum.filter(fn [_, v] -> v == "1" end)
            |> Enum.count()
            |> then(& &1 != 0)

          _ -> false
        end

      _ -> false
    end
  end

  def compilable?() do
    base_root = :crypto.strong_rand_bytes(10) |> Base.encode32(case: :lower)
    base = base_root <> ".c"

    File.write(
      Path.join("/tmp", base), 
      """
      #include <arm_sme.h>

      __arm_locally_streaming
      __arm_new("za")
      void test_arm_new(void) {}
      int main(int argc, char *argv[])
      {
        test_arm_new();
      }
      """
    )
    |> case do
      :ok -> 
        execute("/usr/bin/clang", 
          [
            "-O2", 
            "-Werror",
            "-Wall",
            "-march=armv9-a+sme", 
            base,
            "-o",
            base_root
          ],
          cd: "/tmp",
          stderr_to_stdout: true
        )
        |> case do
          {_, 0} -> true
          _ -> false
        end
      
      {:error, _reason} -> false
    end
  end

  defp execute(executable, options, opts \\ []) do
    System.find_executable(executable)
    |> case do
      nil -> 
        executable
        |> Path.basename()
        |> System.find_executable()
        |> case do
          nil -> false
          executable -> System.cmd(executable, options, opts)
        end

      executable -> System.cmd(executable, options, opts)
    end
  end
end

P.S. このプログラムを踏まえて，Autoconfex を作成しました．

Mix.install([:autoconfex])

defmodule SME do
  def available?() do
    {:ok, pid} = Task.Supervisor.start_link()

    task1 = Task.Supervisor.async(pid, fn -> runnable?() end)
    task2 = Task.Supervisor.async(pid, fn -> compilable?() end)
    
    Task.await(task1) and Task.await(task2)
  end

  def runnable?() do
    case :os.type() do
      {:unix, :darwin} -> 
        case execute("sysctl", ["hw.optional.arm"]) do
          {result, 0} -> 
            result
            |> String.split("\n")
            |> Enum.map(&String.trim/1)
            |> Enum.filter(&String.match?(&1, ~r/FEAT\_SME/))
            |> Enum.map(&String.split(&1, " "))
            |> Enum.filter(fn [_, v] -> v == "1" end)
            |> Enum.count()
            |> then(& &1 != 0)

          _ -> false
        end

      _ -> false
    end
  end

  def compilable?() do
    Autoconfex.compilable_by_cc?(
      "/usr/bin/clang", 
      """
      #include <arm_sme.h>

      __arm_locally_streaming
      __arm_new("za")
      void test_arm_new(void) {}
      int main(int argc, char *argv[])
      {
        test_arm_new();
      }
      """,
      ["-O2", "-march=armv9-a+sme"]
    )
  end

  defp execute(executable, options, opts \\ []) do
    System.find_executable(executable)
    |> case do
      nil -> 
        executable
        |> Path.basename()
        |> System.find_executable()
        |> case do
          nil -> false
          executable -> System.cmd(executable, options, opts)
        end

      executable -> System.cmd(executable, options, opts)
    end
  end
end