amx.hpp

Go to the documentation of this file.

 
#ifndef NUMERICAL_AMX_HPP
#define NUMERICAL_AMX_HPP
 
#include <algorithm>
#include <array>
#include <cstring>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <initializer_list>
 
#include "simd.hpp"
 
#if defined(__APPLE__)
#  include <TargetConditionals.h>
#endif
#if defined(__APPLE__) && defined(__aarch64__) && TARGET_OS_OSX && !defined(NUMERICAL_NO_AMX)
#  define NUMERICAL_AMX_COMPILED // the AMX path is compiled in (tested with #ifdef only)
#  include <sys/sysctl.h>
#endif
 
namespace numerical::detail {
 
#ifdef NUMERICAL_AMX_COMPILED
  enum class amx_op : unsigned
  {
    ldx   = 0,  
    ldy   = 1,  
    ldz   = 4,  
    stz   = 5,  
    fma64 = 10, 
    fma32 = 12, 
  };
 
  template <amx_op Op>
  inline void amx(std::uint64_t operand) noexcept
  {
    __asm__ volatile (".word (0x201000 + (%0 << 5) + 0%1 - ((0%1 >> 4) * 6))" // NOLINT(hicpp-no-assembler): no compiler emits the AMX
                      :
                      : "i" (static_cast<unsigned>(Op)), "r" (operand)
                      : "memory");
  }
 
  template <bool Off>
  inline void amx_switch() noexcept
  {
    __asm__ volatile ("nop\nnop\nnop\n.word (0x201000 + (17 << 5) + %0)" // NOLINT(hicpp-no-assembler): no compiler emits the AMX
                      :
                      : "i" (static_cast<unsigned>(Off))
                      : "memory");
  }
 
  class amx_session
  {
  public:
 
    amx_session() noexcept
    {
      amx_switch<false>();
    }
 
    amx_session(const amx_session&)            = delete;
    amx_session& operator=(const amx_session&) = delete;
    amx_session(amx_session&&)                 = delete;
    amx_session& operator=(amx_session&&)      = delete;
 
    ~amx_session()
    {
      amx_switch<true>();
    }
  };
 
  [[nodiscard]] inline std::uint64_t amx_address(const void  * at,
                                                 std::uint64_t index) noexcept
  {
    return (index << 56U) | reinterpret_cast<std::uintptr_t>(at); // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast): the instruction takes the address as an integer
  }
 
  inline constexpr std::size_t amx_dot_group {8};
 
  inline void amx_dot_rows_f64(const double* m,
                               std::size_t   ld,
                               std::size_t   rows,
                               const double* x,
                               std::size_t   count,
                               double*       out) noexcept
  {
    alignas(64) static constexpr std::array<double, 8> zero            {};
    alignas(64) std::array<std::array<double, 8>, amx_dot_group> lanes {};
    const std::size_t whole                                            {count - (count % 8)};
    for (std::size_t r0 {0}; r0 < rows; r0 += amx_dot_group) {
      const std::size_t group {std::min(amx_dot_group, rows - r0)};
      for (std::size_t r {0}; r < group; ++r) {
        amx<amx_op::ldz>(amx_address(zero.data(), r));
      }
      for (std::size_t j {0}; j < whole; j += 8) {
        amx<amx_op::ldx>(amx_address(x + j, 0));
        for (std::size_t r {0}; r < group; ++r) {
          amx<amx_op::ldy>(amx_address(m + ((r0 + r) * ld) + j, r));
          amx<amx_op::fma64>((std::uint64_t {1} << 63U) | (std::uint64_t {r} << 20U) | (std::uint64_t {r} *64U)); // vector mode: Z row r += X0 ∘ Y r
        }
      }
      for (std::size_t r {0}; r < group; ++r) { // every row stored before any is read: a read right behind its store waited on it
        amx<amx_op::stz>(amx_address(lanes[r].data(), r));
      }
      for (std::size_t r {0}; r < group; ++r) {
        for (std::size_t j {whole}; j < count; ++j) {
          lanes[r][j % 8] = std::fma(m[((r0 + r) * ld) + j], x[j], lanes[r][j % 8]);
        }
        double sum {lanes[r][0]};
        for (std::size_t l {1}; l < 8; ++l) {
          sum += lanes[r][l];
        }
        out[r0 + r] = sum;
      }
    }
  }
 
  inline constexpr std::size_t amx_mr {32};
 
  inline constexpr std::size_t amx_nr {16};
 
  inline constexpr std::size_t amx_least_extent {16};
 
  inline constexpr std::size_t amx_least_work {std::size_t {1} << 16U};
 
  [[nodiscard]] constexpr bool amx_suits(std::size_t m,
                                         std::size_t n,
                                         std::size_t k) noexcept
  {
    return (m >= amx_least_extent && n >= amx_least_extent) || m * n * k >= amx_least_work;
  }
 
  inline constexpr std::size_t amx_f32_side {32};
 
  inline constexpr std::size_t amx_f32_least_width {28};
 
  [[nodiscard]] constexpr bool amx_suits_f32(std::size_t m,
                                             std::size_t n,
                                             std::size_t k) noexcept
  {
    const bool whole_tiles   {n % amx_f32_side == 0};
    const bool wide          {m >= amx_least_extent && n >= amx_f32_least_width && m * n * k >= amx_least_work};
    const bool tall_and_deep {m >= 64 && k >= 64};
    return amx_suits(m, n, k) && (whole_tiles || wide || tall_and_deep);
  }
 
  inline constexpr std::size_t amx_in_place_stride {512};
 
  [[gnu::noinline, gnu::aligned(64)]] inline void amx_depth_f64(std::size_t   kc,
                                                                const double* ap,
                                                                const double* bp,
                                                                std::size_t   ldb) noexcept
  {
    for (std::size_t p {0}; p < kc; ++p) {
      const double* const a {ap + (p * amx_mr)};
      const double* const b {bp + (p * ldb)};
      for (std::uint64_t ya {0}; ya < 4; ++ya) {
        amx<amx_op::ldy>(amx_address(a + (ya * 8), ya));
      }
      amx<amx_op::ldx>(amx_address(b, 0));
      amx<amx_op::ldx>(amx_address(b + 8, 1));
      for (std::uint64_t ya {0}; ya < 4; ++ya) {
        for (std::uint64_t xb {0}; xb < 2; ++xb) {
          // Matrix mode: Z row (2·ya + xb), X register xb (64 bytes each), Y register ya.
          amx<amx_op::fma64>((((ya * 2) + xb) << 20U) | ((xb * 64) << 10U) | (ya * 64));
        }
      }
    }
  }
 
  inline void amx_full_tile_f64(std::size_t   kc,
                                const double* ap,
                                const double* bp,
                                std::size_t   ldb,
                                double*       tile,
                                std::size_t   ld) noexcept
  {
    for (std::uint64_t ya {0}; ya < 4; ++ya) {
      for (std::uint64_t xb {0}; xb < 2; ++xb) {
        for (std::uint64_t j {0}; j < 8; ++j) {
          amx<amx_op::ldz>(amx_address(tile + (((ya * 8) + j) * ld) + (xb * 8), (j * 8) + (ya * 2) + xb));
        }
      }
    }
    amx_depth_f64(kc, ap, bp, ldb);
    for (std::uint64_t ya {0}; ya < 4; ++ya) {
      for (std::uint64_t xb {0}; xb < 2; ++xb) {
        for (std::uint64_t j {0}; j < 8; ++j) {
          amx<amx_op::stz>(amx_address(tile + (((ya * 8) + j) * ld) + (xb * 8), (j * 8) + (ya * 2) + xb));
        }
      }
    }
  }
 
  [[gnu::noinline, gnu::aligned(64)]] inline void amx_depth_f32(std::size_t  kc,
                                                                const float* ap,
                                                                const float* bp,
                                                                std::size_t  ldb) noexcept
  {
    for (std::size_t p {0}; p < kc; ++p) {
      const float* const a {ap + (p * amx_f32_side)};
      const float* const b {bp + (p * ldb)};
      amx<amx_op::ldy>(amx_address(a, 0));
      amx<amx_op::ldy>(amx_address(a + 16, 1));
      amx<amx_op::ldx>(amx_address(b, 0));
      amx<amx_op::ldx>(amx_address(b + 16, 1));
      for (std::uint64_t ya {0}; ya < 2; ++ya) {
        for (std::uint64_t xb {0}; xb < 2; ++xb) {
          amx<amx_op::fma32>((((ya * 2) + xb) << 20U) | ((xb * 64) << 10U) | (ya * 64));
        }
      }
    }
  }
 
  inline void amx_depth(std::size_t   kc,
                        const double* ap,
                        const double* bp,
                        std::size_t   ldb) noexcept
  {
    amx_depth_f64(kc, ap, bp, ldb);
  }
 
  inline void amx_depth(std::size_t  kc,
                        const float* ap,
                        const float* bp,
                        std::size_t  ldb) noexcept
  {
    amx_depth_f32(kc, ap, bp, ldb);
  }
 
  template <typename T>
  inline constexpr amx_op amx_multiply_add {amx_op::fma32};
 
  template <>
  inline constexpr amx_op amx_multiply_add<double> {amx_op::fma64};
 
  template <typename T,
            std::uint64_t Bands,
            std::uint64_t Halves>
  [[gnu::noinline, gnu::aligned(64)]] void amx_depth_part(std::size_t kc,
                                                          const T*    ap,
                                                          const T*    bp,
                                                          std::size_t ldb) noexcept
  {
    constexpr std::size_t lanes        {64 / sizeof(T)}; // a register's elements
    for (std::size_t p {0}; p < kc; ++p) {
      const T* const a {ap + (p * amx_mr)};
      const T* const b {bp + (p * ldb)};
      for (std::uint64_t ya {0}; ya < Bands; ++ya) {
        amx<amx_op::ldy>(amx_address(a + (ya * lanes), ya));
      }
      for (std::uint64_t xb {0}; xb < Halves; ++xb) {
        amx<amx_op::ldx>(amx_address(b + (xb * lanes), xb));
      }
      for (std::uint64_t ya {0}; ya < Bands; ++ya) {
        for (std::uint64_t xb {0}; xb < Halves; ++xb) {
          amx<amx_multiply_add<T>>((((ya * 2) + xb) << 20U) | ((xb * 64) << 10U) | (ya * 64));
        }
      }
    }
  }
 
  template <typename T,
            std::uint64_t Bands>
  void amx_depth_part_halves(std::size_t   kc,
                             const T*      ap,
                             const T*      bp,
                             std::size_t   ldb,
                             std::uint64_t halves) noexcept
  {
    if (halves == 1) {
      amx_depth_part<T, Bands, 1>(kc, ap, bp, ldb);
    }
    else {
      amx_depth_part<T, Bands, 2>(kc, ap, bp, ldb);
    }
  }
 
  inline void amx_depth_part(std::size_t   kc,
                             const double* ap,
                             const double* bp,
                             std::size_t   ldb,
                             std::uint64_t bands,
                             std::uint64_t halves) noexcept
  {
    if (bands == 1) {
      amx_depth_part_halves<double, 1>(kc, ap, bp, ldb, halves);
    }
    else if (bands == 2) {
      amx_depth_part_halves<double, 2>(kc, ap, bp, ldb, halves);
    }
    else if (bands == 3) {
      amx_depth_part_halves<double, 3>(kc, ap, bp, ldb, halves);
    }
    else { // every band: a partial tile, so one half (the whole tile goes to amx_depth)
      amx_depth_part<double, 4, 1>(kc, ap, bp, ldb);
    }
  }
 
  inline void amx_depth_part(std::size_t   kc,
                             const float*  ap,
                             const float*  bp,
                             std::size_t   ldb,
                             std::uint64_t bands,
                             std::uint64_t halves) noexcept
  {
    if (bands == 1) {
      amx_depth_part_halves<float, 1>(kc, ap, bp, ldb, halves);
    }
    else { // every band: a partial tile, so one half (the whole tile goes to amx_depth)
      amx_depth_part<float, 2, 1>(kc, ap, bp, ldb);
    }
  }
 
  template <typename T>
  void amx_partial_tile(std::size_t kc,
                        const T*    ap,
                        const T*    bp,
                        std::size_t ldb,
                        T*          c,
                        std::size_t ldc,
                        std::size_t rows,
                        std::size_t cols) noexcept
  {
    static_assert(amx_f32_side == amx_mr, "both tiles have amx_mr rows");
    constexpr std::size_t           lanes {64 / sizeof(T)};                // a half row: one Z row
    constexpr std::uint64_t         bands {amx_mr / lanes};                // accumulators down the tile
    alignas(64) std::array<T, amx_mr * lanes> cut {};                      // the half row the last column cuts, per row
    const std::uint64_t             whole_halves {cols / lanes};           // halves wholly inside C
    const std::size_t               cut_width {cols % lanes};              // the cut half's elements inside C
    const auto                      each_half = [&](auto&& at) {
                                                  for (std::uint64_t ya {0}; ya < bands; ++ya) {
                                                    for (std::uint64_t j {0}; j < lanes; ++j) {
                                                      const std::size_t row {(ya * lanes) + j};
                                                      if (row >= rows) {
                                                        continue;
                                                      }
                                                      for (std::uint64_t xb {0}; xb < 2; ++xb) {
                                                        const std::uint64_t z {(j * 2 * bands) + (ya * 2) + xb};
                                                        if (xb < whole_halves) {
                                                          at(c + (row * ldc) + (xb * lanes), z);
                                                        }
                                                        else if (xb == whole_halves && cut_width != 0) {
                                                          at(cut.data() + (row * lanes), z);
                                                        }
                                                      }
                                                    }
                                                  }
                                                };
    if (cut_width != 0) {
      for (std::size_t row {0}; row < rows; ++row) {
        for (std::size_t e {0}; e < cut_width; ++e) {
          cut[(row * lanes) + e] = c[(row * ldc) + (whole_halves * lanes) + e];
        }
      }
    }
    each_half([](T* half, std::uint64_t z) { amx<amx_op::ldz>(amx_address(half, z)); });
    const std::uint64_t used_bands  {(rows + lanes - 1) / lanes};                          // the accumulators C's rows fill
    const std::uint64_t used_halves {std::min<std::uint64_t>(2, (cols + lanes - 1) / lanes)};
    if (used_bands == bands && used_halves == 2) {
      amx_depth(kc, ap, bp, ldb);
    }
    else {
      amx_depth_part(kc, ap, bp, ldb, used_bands, used_halves);
    }
    each_half([](T* half, std::uint64_t z) { amx<amx_op::stz>(amx_address(half, z)); });
    if (cut_width != 0) {
      for (std::size_t row {0}; row < rows; ++row) {
        for (std::size_t e {0}; e < cut_width; ++e) {
          c[(row * ldc) + (whole_halves * lanes) + e] = cut[(row * lanes) + e];
        }
      }
    }
  }
 
  inline void amx_full_tile_f32(std::size_t  kc,
                                const float* ap,
                                const float* bp,
                                std::size_t  ldb,
                                float*       tile,
                                std::size_t  ld) noexcept
  {
    for (std::uint64_t ya {0}; ya < 2; ++ya) {
      for (std::uint64_t xb {0}; xb < 2; ++xb) {
        for (std::uint64_t j {0}; j < 16; ++j) {
          amx<amx_op::ldz>(amx_address(tile + (((ya * 16) + j) * ld) + (xb * 16), (j * 4) + (ya * 2) + xb));
        }
      }
    }
    amx_depth_f32(kc, ap, bp, ldb);
    for (std::uint64_t ya {0}; ya < 2; ++ya) {
      for (std::uint64_t xb {0}; xb < 2; ++xb) {
        for (std::uint64_t j {0}; j < 16; ++j) {
          amx<amx_op::stz>(amx_address(tile + (((ya * 16) + j) * ld) + (xb * 16), (j * 4) + (ya * 2) + xb));
        }
      }
    }
  }
 
  inline void amx_tile_f32(std::size_t  kc,
                           const float* ap,
                           const float* bp,
                           std::size_t  ldb,
                           float*       c,
                           std::size_t  ldc,
                           std::size_t  rows,
                           std::size_t  cols) noexcept
  {
    if (rows == amx_f32_side && cols == amx_f32_side) {
      amx_full_tile_f32(kc, ap, bp, ldb, c, ldc);
      return;
    }
    amx_partial_tile(kc, ap, bp, ldb, c, ldc, rows, cols);
  }
 
  inline void amx_tile_f64(std::size_t   kc,
                           const double* ap,
                           const double* bp,
                           std::size_t   ldb,
                           double*       c,
                           std::size_t   ldc,
                           std::size_t   rows,
                           std::size_t   cols) noexcept
  {
    if ((rows < amx_mr || cols < amx_nr) && (cols < 8 || cols % 8 == 0)) {
      amx_partial_tile(kc, ap, bp, ldb, c, ldc, rows, cols);
      return;
    }
    through_edge<amx_mr, amx_nr>(c, ldc, rows, cols, [&](double* tile, std::size_t ld) noexcept { amx_full_tile_f64(kc, ap, bp, ldb, tile, ld); });
  }
 
  inline constexpr std::size_t amx_stream_least {2048};
 
  inline void amx_pack_a(const double* a,
                         std::size_t   lda,
                         std::size_t   rows,
                         std::size_t   kc,
                         double*       out) noexcept
  {
    using pair = double __attribute__((vector_size(16)));
    std::size_t ir {0};
    for (; rows * kc >= amx_stream_least && ir + amx_mr <= rows; ir += amx_mr) {
      double* const sliver {out + (ir * kc)};
      std::size_t   p      {0};
      for (; p + 2 <= kc; p += 2) { // two columns of four rows a step: the rows' pairs read whole, their lanes swapped
        for (std::size_t i {0}; i < amx_mr; i += 4) {
          std::array<pair, 4> row {};
          for (std::size_t r {0}; r < 4; ++r) {
            std::memcpy(&row.at(r), a + ((ir + i + r) * lda) + p, sizeof(pair));
          }
          const pair here_low  {__builtin_shufflevector(row[0], row[1], 0, 2)};
          const pair here_high {__builtin_shufflevector(row[2], row[3], 0, 2)};
          const pair next_low  {__builtin_shufflevector(row[0], row[1], 1, 3)};
          const pair next_high {__builtin_shufflevector(row[2], row[3], 1, 3)};
          __asm__ volatile ("stnp %q0, %q1, [%2]" // NOLINT(hicpp-no-assembler): as below, four elements at once
                            :
                            : "w" (here_low), "w" (here_high), "r" (sliver + (p * amx_mr) + i)
                            : "memory");
          __asm__ volatile ("stnp %q0, %q1, [%2]" // NOLINT(hicpp-no-assembler): as below, four elements at once
                            :
                            : "w" (next_low), "w" (next_high), "r" (sliver + ((p + 1) * amx_mr) + i)
                            : "memory");
        }
      }
      for (; p < kc; ++p) {
        for (std::size_t i {0}; i < amx_mr; i += 2) {
          const double first  {a[((ir + i) * lda) + p]};
          const double second {a[((ir + i + 1) * lda) + p]};
          __asm__ volatile ("stnp %d0, %d1, [%2]" // NOLINT(hicpp-no-assembler): no builtin emits a pair past the L1 under both compilers
                            :
                            : "w" (first), "w" (second), "r" (sliver + (p * amx_mr) + i)
                            : "memory");
        }
      }
    }
    for (; ir + amx_mr <= rows; ir += amx_mr) { // whole slivers below the threshold, the loop free of the padding's test
      double* const sliver {out + (ir * kc)};
      std::size_t   p      {0};
      for (; p + 2 <= kc; p += 2) { // as above, by ordinary stores
        for (std::size_t i {0}; i < amx_mr; i += 2) {
          pair first;
          pair second;
          std::memcpy(&first, a + ((ir + i) * lda) + p, sizeof(pair));
          std::memcpy(&second, a + ((ir + i + 1) * lda) + p, sizeof(pair));
          const pair here {__builtin_shufflevector(first, second, 0, 2)};
          const pair next {__builtin_shufflevector(first, second, 1, 3)};
          std::memcpy(sliver + (p * amx_mr) + i, &here, sizeof(pair));
          std::memcpy(sliver + ((p + 1) * amx_mr) + i, &next, sizeof(pair));
        }
      }
      for (; p < kc; ++p) {
        for (std::size_t i {0}; i < amx_mr; ++i) {
          sliver[(p * amx_mr) + i] = a[((ir + i) * lda) + p];
        }
      }
    }
    if (ir < rows) {
      double* const sliver {out + (ir * kc)};
      for (std::size_t p {0}; p < kc; ++p) {
        for (std::size_t i {0}; i < amx_mr; ++i) {
          sliver[(p * amx_mr) + i] = ir + i < rows ? a[((ir + i) * lda) + p] : 0.0;
        }
      }
    }
  }
 
  inline void amx_pack_a_f32(const float* a,
                             std::size_t  lda,
                             std::size_t  rows,
                             std::size_t  kc,
                             float*       out) noexcept
  {
    using quad = float __attribute__((vector_size(16)));
    constexpr std::size_t stream_least {amx_stream_least * (sizeof(double) / sizeof(float))};
    // Four rows' four elements as four columns.
    const auto transposed = [](const quad& r0, const quad& r1, const quad& r2, const quad& r3) noexcept {
                              const quad low01  {__builtin_shufflevector(r0, r1, 0, 4, 1, 5)};
                              const quad high01 {__builtin_shufflevector(r0, r1, 2, 6, 3, 7)};
                              const quad low23  {__builtin_shufflevector(r2, r3, 0, 4, 1, 5)};
                              const quad high23 {__builtin_shufflevector(r2, r3, 2, 6, 3, 7)};
                              return std::array<quad, 4> {__builtin_shufflevector(low01, low23, 0, 1, 4, 5), __builtin_shufflevector(low01, low23, 2, 3, 6, 7),
                                                          __builtin_shufflevector(high01, high23, 0, 1, 4, 5), __builtin_shufflevector(high01, high23, 2, 3, 6, 7)};
                            };
    // A column's eight rows, past the L1.
    const auto store = [](const quad& upper, const quad& lower, float* at) noexcept {
                         __asm__ volatile ("stnp %q0, %q1, [%2]" // NOLINT(hicpp-no-assembler): as amx_pack_a, eight elements at once
                                           :
                                           : "w" (upper), "w" (lower), "r" (at)
                                           : "memory");
                       };
    std::size_t           ir           {0};
    for (; rows * kc >= stream_least && ir + amx_f32_side <= rows; ir += amx_f32_side) {
      float* const sliver {out + (ir * kc)};
      std::size_t  p      {0};
      for (; p + 4 <= kc; p += 4) { // four columns of eight rows a step: two blocks of 4 × 4 transposed
        for (std::size_t i {0}; i < amx_f32_side; i += 8) {
          std::array<quad, 8> row {};
          for (std::size_t r {0}; r < 8; ++r) {
            std::memcpy(&row.at(r), a + ((ir + i + r) * lda) + p, sizeof(quad));
          }
          const std::array<quad, 4> upper {transposed(row[0], row[1], row[2], row[3])};
          const std::array<quad, 4> lower {transposed(row[4], row[5], row[6], row[7])};
          float* const              at    {sliver + (p * amx_f32_side) + i};
          store(upper[0], lower[0], at);
          store(upper[1], lower[1], at + amx_f32_side);
          store(upper[2], lower[2], at + (2 * amx_f32_side));
          store(upper[3], lower[3], at + (3 * amx_f32_side));
        }
      }
      for (; p < kc; ++p) {
        for (std::size_t i {0}; i < amx_f32_side; i += 2) {
          const float first  {a[((ir + i) * lda) + p]};
          const float second {a[((ir + i + 1) * lda) + p]};
          __asm__ volatile ("stnp %s0, %s1, [%2]" // NOLINT(hicpp-no-assembler): as amx_pack_a, a pair past the L1
                            :
                            : "w" (first), "w" (second), "r" (sliver + (p * amx_f32_side) + i)
                            : "memory");
        }
      }
    }
    for (; ir < rows; ir += amx_f32_side) {
      float* const      sliver {out + (ir * kc)};
      const std::size_t height {std::min(amx_f32_side, rows - ir)};
      for (std::size_t p {0}; p < kc; ++p) {
        for (std::size_t i {0}; i < amx_f32_side; ++i) {
          sliver[(p * amx_f32_side) + i] = i < height ? a[((ir + i) * lda) + p] : 0.0F;
        }
      }
    }
  }
 
  [[nodiscard]] constexpr bool amx_family_known(std::uint32_t family) noexcept
  {
    switch (family) {
      case 0x1b588bb3U: // FIRESTORM_ICESTORM: M1, M1 Pro, Max, Ultra
      case 0xda33d83dU: // BLIZZARD_AVALANCHE: M2 family
      case 0xfa33415eU: // IBIZA: M3
      case 0x5f4dea93U: // LOBOS: M3 Pro
      case 0x72015832U: // PALMA: M3 Max
      case 0x6f5129acU: // DONAN: M4
      case 0x17d5b93aU: // BRAVA: M4 Pro, Max
        return true;
      default:
        return false;
    }
  }
 
  [[nodiscard]] inline std::uint32_t cpu_family() noexcept
  {
    std::uint32_t family {0};
    std::size_t   size   {sizeof family};
    (void)sysctlbyname("hw.cpufamily", &family, &size, nullptr, 0); // on failure family stays 0
    return family;
  }
 
  [[nodiscard]] inline bool amx_dot_self_test() noexcept
  {
    constexpr std::size_t                    rows  {11};
    constexpr std::size_t                    count {21};
    std::array<double, rows * count>         m     {};
    std::array<double, count>                x     {};
    std::array<double, rows>                 got   {};
    std::array<double, rows>                 want  {};
    for (std::size_t i {0}; i < m.size(); ++i) {
      m[i] = std::ldexp(static_cast<double>((i * 37) % 101) - 50.0, -static_cast<int>(i % 9));
    }
    for (std::size_t j {0}; j < count; ++j) {
      x[j] = 1.0 / static_cast<double>((j % 13) + 3);
    }
    for (std::size_t r {0}; r < rows; ++r) {
      std::array<double, 8> lane {};
      for (std::size_t j {0}; j < count; ++j) {
        lane[j % 8] = std::fma(m[(r * count) + j], x[j], lane[j % 8]);
      }
      double sum {lane[0]};
      for (std::size_t l {1}; l < 8; ++l) {
        sum += lane[l];
      }
      want[r] = sum;
    }
    {
      const amx_session on {};
      amx_dot_rows_f64(m.data(), count, rows, x.data(), count, got.data());
    }
    return got == want;
  }
 
  [[nodiscard]] inline bool amx_f32_self_test() noexcept
  {
    constexpr std::size_t                          depth {7};
    constexpr std::size_t                          side  {amx_f32_side};
    std::array<float, depth * side>                a     {};
    std::array<float, depth * side>                b     {};
    std::array<float, side * side>                 c     {};
    std::array<float, side * side>                 want  {};
    for (std::size_t i {0}; i < depth * side; ++i) {
      a[i] = std::ldexp(static_cast<float>((i * 37) % 101) - 50.0F, -static_cast<int>(i % 9));
      b[i] = 1.0F / static_cast<float>((i % 13) + 3);
    }
    for (std::size_t i {0}; i < side * side; ++i) {
      c[i]    = static_cast<float>(i % 17) * 0.1F;
      want[i] = c[i];
    }
    for (const std::size_t rows : {side, std::size_t {19}}) {
      const std::size_t cols {rows == side ? side : std::size_t {11}};
      for (std::size_t i {0}; i < rows; ++i) {
        for (std::size_t j {0}; j < cols; ++j) {
          for (std::size_t p {0}; p < depth; ++p) {
            want[(i * side) + j] = std::fma(a[(p * side) + i], b[(p * side) + j], want[(i * side) + j]);
          }
        }
      }
      {
        const amx_session on {};
        amx_tile_f32(depth, a.data(), b.data(), side, c.data(), side, rows, cols);
      }
    }
    return c == want;
  }
 
  [[nodiscard]] inline bool amx_self_test() noexcept
  {
    constexpr std::size_t               depth                 {7};
    std::array<double, depth * amx_mr>  a                     {};
    std::array<double, depth * amx_nr>  b                     {};
    std::array<double, amx_mr * amx_nr> c                     {};
    std::array<double, amx_mr * amx_nr> want                  {};
    for (std::size_t i {0}; i < depth * amx_mr; ++i) {
      a[i] = std::ldexp(static_cast<double>((i * 37) % 101) - 50.0, -static_cast<int>(i % 9));
    }
    for (std::size_t i {0}; i < depth * amx_nr; ++i) {
      b[i] = 1.0 / static_cast<double>((i % 13) + 3);
    }
    for (std::size_t i {0}; i < amx_mr * amx_nr; ++i) {
      c[i]    = static_cast<double>(i % 17) * 0.1;
      want[i] = c[i];
    }
    for (const std::size_t rows : {amx_mr, std::size_t {19}}) {
      const std::size_t cols {rows == amx_mr ? amx_nr : std::size_t {11}};
      for (std::size_t i {0}; i < rows; ++i) {
        for (std::size_t j {0}; j < cols; ++j) {
          for (std::size_t p {0}; p < depth; ++p) {
            want[(i * amx_nr) + j] = std::fma(a[(p * amx_mr) + i], b[(p * amx_nr) + j], want[(i * amx_nr) + j]);
          }
        }
      }
      {
        const amx_session on {};
        amx_tile_f64(depth, a.data(), b.data(), amx_nr, c.data(), amx_nr, rows, cols);
      }
    }
    const std::array<bool, 3> passed {c == want, amx_f32_self_test(), amx_dot_self_test()};
    return std::count(passed.begin(), passed.end(), true) == 3; // no branch on a check that only a fault fails
  }
 
  [[nodiscard]] inline bool running_virtualised() noexcept
  {
    int         present {0};
    std::size_t size    {sizeof present};
    (void)sysctlbyname("kern.hv_vmm_present", &present, &size, nullptr, 0); // on failure present stays 0
    return present != 0;
  }
 
  [[nodiscard]] inline std::atomic<bool>& amx_probed() noexcept
  {
    static std::atomic<bool> probed {false};
    return probed;
  }
 
  [[nodiscard]] inline bool amx_usable_on(std::uint32_t family,
                                          bool          virtualised) noexcept
  {
    if (!amx_family_known(family) || virtualised) {
      return false;
    }
    amx_probed().store(true);
    return amx_self_test();
  }
 
#endif
 
  [[nodiscard]] inline bool amx_usable() noexcept
  {
#ifdef NUMERICAL_AMX_COMPILED
    static const bool usable {amx_usable_on(cpu_family(), running_virtualised())};
    return usable;
#else
    return false;
#endif
  }
} // namespace numerical::detail
 
#endif // NUMERICAL_AMX_HPP