native.arm.bf16.ccm

// SPDX-License-Identifier: BSD-2-Clause OR Apache-2.0
module;
#include "native/isa_import.h"
#include "native/arm/bf16.h"
#include "native/arm/bf16_constexpr.h"
export module native.arm.bf16;
export import native.arm.features;
export import native.simd;
#if NATIVE_HOST_NEON || defined(NATIVE_DOXYGEN)
namespace native::detail {
  template<int Lane,std::size_t N,std::size_t M,isa<arm> Arch>
  consteval simd<float,N,Arch> bfdot_value(simd<float,N,Arch> acc,
      simd<bf16,2*N,Arch> a,simd<bf16,M,Arch> b) noexcept {
    std::array<std::uint32_t,N> c{};
    std::array<std::uint16_t,2*N> av{};
    std::array<std::uint16_t,M> bv{};
    acc.store_bits(c.data());a.store_bits(av.data());b.store_bits(bv.data());
    for(std::size_t i=0;i<N;++i) {
      auto j=Lane<0?2*i:2*Lane;
      c[i]=arm_bfdot_bits(c[i],av[2*i],av[2*i+1],bv[j],bv[j+1]);
    }
    return simd<float,N,Arch>::load_bits(c.data());
  }
  template<isa<arm> Arch>
  consteval simd<float,4,Arch> bfmmla_value(simd<float,4,Arch> acc,
      simd<bf16,8,Arch> a,simd<bf16,8,Arch> b) noexcept {
    std::array<std::uint32_t,4> c{};
    std::array<std::uint16_t,8> av{},bv{};
    acc.store_bits(c.data());a.store_bits(av.data());b.store_bits(bv.data());
    for(unsigned row=0;row<2;++row) for(unsigned column=0;column<2;++column)
      for(unsigned pair=0;pair<2;++pair) {
        auto i=4*row+2*pair,j=4*column+2*pair;
        c[2*row+column]=arm_bfdot_bits(c[2*row+column],av[i],av[i+1],bv[j],bv[j+1]);
      }
    return simd<float,4,Arch>::load_bits(c.data());
  }
  template<bool Top,int Lane,std::size_t M,isa<arm> Arch>
  consteval simd<float,4,Arch> bfmlal_value(simd<float,4,Arch> acc,
      simd<bf16,8,Arch> a,simd<bf16,M,Arch> b) noexcept {
    std::array<std::uint32_t,4> c{};
    std::array<std::uint16_t,8> av{};
    std::array<std::uint16_t,M> bv{};
    acc.store_bits(c.data());a.store_bits(av.data());b.store_bits(bv.data());
    for(unsigned i=0;i<4;++i) {
      auto j=2*i+Top;
      c[i]=constexpr_float::fma_bits<constexpr_float::binary32>(
        std::uint32_t(av[j])<<16,std::uint32_t(bv[Lane<0?j:Lane])<<16,c[i]);
    }
    return simd<float,4,Arch>::load_bits(c.data());
  }
}
export namespace native {

  template<isa<arm> Arch>
    requires(Arch.has(arm_feature::neon_bf16))
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 2, Arch> bfdot(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<-1>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot<Arch>(vget_low_f32(acc.to_storage().to_native()), __builtin_bit_cast(bfloat16x4_t, a.to_native()), __builtin_bit_cast(bfloat16x4_t, b.to_native()));
      return simd<float, 2, Arch>::from_storage(
        simd<float, 4, Arch>::from_native(vcombine_f32(result, vdup_n_f32(0.f))));
    }
  }

  template<isa<arm> Arch>
    requires(requires { sizeof(simd<float, 2, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && !(Arch.has(arm_feature::neon_bf16)))
  native_nodiscard consteval
  simd<float, 2, Arch> bfdot(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    return detail::bfdot_value<-1>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 2)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 2, Arch> bfdot_lane(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot_lane<Arch, Lane>(vget_low_f32(acc.to_storage().to_native()), __builtin_bit_cast(bfloat16x4_t, a.to_native()), __builtin_bit_cast(bfloat16x4_t, b.to_native()));
      return simd<float, 2, Arch>::from_storage(
        simd<float, 4, Arch>::from_native(vcombine_f32(result, vdup_n_f32(0.f))));
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 2, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 2)
  native_nodiscard consteval
  simd<float, 2, Arch> bfdot_lane(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    return detail::bfdot_value<Lane>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 4)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 2, Arch> bfdot_lane(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot_lane<Arch, Lane>(vget_low_f32(acc.to_storage().to_native()), __builtin_bit_cast(bfloat16x4_t, a.to_native()), b.to_native());
      return simd<float, 2, Arch>::from_storage(
        simd<float, 4, Arch>::from_native(vcombine_f32(result, vdup_n_f32(0.f))));
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 2, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 4)
  native_nodiscard consteval
  simd<float, 2, Arch> bfdot_lane(simd<float, 2, Arch> acc, simd<bf16, 4, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfdot_value<Lane>(acc,a,b);
  }

  template<isa<arm> Arch>
    requires(Arch.has(arm_feature::neon_bf16))
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfdot(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<-1>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot<Arch>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)))
  native_nodiscard consteval
  simd<float, 4, Arch> bfdot(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfdot_value<-1>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 2)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfdot_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot_lane<Arch, Lane>(acc.to_native(), a.to_native(), __builtin_bit_cast(bfloat16x4_t, b.to_native()));
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 2)
  native_nodiscard consteval
  simd<float, 4, Arch> bfdot_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    return detail::bfdot_value<Lane>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 4)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfdot_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfdot_value<Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfdot_lane<Arch, Lane>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 4)
  native_nodiscard consteval
  simd<float, 4, Arch> bfdot_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfdot_value<Lane>(acc,a,b);
  }

  template<isa<arm> Arch>
    requires(Arch.has(arm_feature::neon_bf16))
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmmla(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfmmla_value(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmmla<Arch>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)))
  native_nodiscard consteval
  simd<float, 4, Arch> bfmmla(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfmmla_value(acc,a,b);
  }

  template<isa<arm> Arch>
    requires(Arch.has(arm_feature::neon_bf16))
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalb(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<false,-1>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalb<Arch>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)))
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalb(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfmlal_value<false,-1>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 4)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalb_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<false,Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalb_lane<Arch, Lane>(acc.to_native(), a.to_native(), __builtin_bit_cast(bfloat16x4_t, b.to_native()));
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 4)
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalb_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    return detail::bfmlal_value<false,Lane>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 8)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalb_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<false,Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalb_lane<Arch, Lane>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 8)
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalb_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfmlal_value<false,Lane>(acc,a,b);
  }

  template<isa<arm> Arch>
    requires(Arch.has(arm_feature::neon_bf16))
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalt(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<true,-1>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalt<Arch>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)))
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalt(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfmlal_value<true,-1>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 4)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalt_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<true,Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalt_lane<Arch, Lane>(acc.to_native(), a.to_native(), __builtin_bit_cast(bfloat16x4_t, b.to_native()));
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && requires { sizeof(simd<bf16, 4, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 4)
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalt_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 4, Arch> b) noexcept {
    return detail::bfmlal_value<true,Lane>(acc,a,b);
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(Arch.has(arm_feature::neon_bf16) && Lane < 8)
  native_nodiscard native_inline __attribute__((target("bf16")))
  constexpr simd<float, 4, Arch> bfmlalt_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    if consteval { return detail::bfmlal_value<true,Lane>(acc,a,b); } else {
      auto result = detail::arm_bf16::bfmlalt_lane<Arch, Lane>(acc.to_native(), a.to_native(), b.to_native());
      return simd<float, 4, Arch>::from_native(result);
    }
  }

  template<isa<arm> Arch, unsigned Lane>
    requires(requires { sizeof(simd<float, 4, Arch>); } && requires { sizeof(simd<bf16, 8, Arch>); } && !(Arch.has(arm_feature::neon_bf16)) && Lane < 8)
  native_nodiscard consteval
  simd<float, 4, Arch> bfmlalt_lane(simd<float, 4, Arch> acc, simd<bf16, 8, Arch> a, simd<bf16, 8, Arch> b) noexcept {
    return detail::bfmlal_value<true,Lane>(acc,a,b);
  }
 
  // Exact deduction rejects unrelated vectors and invalid immediates before
  // Clang's lax vector conversions can select an overload for another shape.
  template<isa<arm> Arch, class R, class A, class B> void bfdot(R, A, B) = delete;
  template<isa<arm> Arch, class R, class A, class B> void bfmmla(R, A, B) = delete;
  template<isa<arm> Arch, class R, class A, class B> void bfmlalb(R, A, B) = delete;
  template<isa<arm> Arch, class R, class A, class B> void bfmlalt(R, A, B) = delete;
  template<isa<arm> Arch, unsigned Lane, class R, class A, class B> void bfdot_lane(R, A, B) = delete;
  template<isa<arm> Arch, unsigned Lane, class R, class A, class B> void bfmlalb_lane(R, A, B) = delete;
  template<isa<arm> Arch, unsigned Lane, class R, class A, class B> void bfmlalt_lane(R, A, B) = delete;

}
#endif