1use super::{tile_loop_generic, validate_gemm_sizes, TiledGemm};
2#[cfg(target_arch = "x86_64")]
3use crate::cpu::{AmxSupport, Avx512Support};
4use eunomia::{Bf16, F32};
5use hermes_simd_core::view::{SimdError, TileMatrixMultiply};
6use hermes_simd_intrinsics::Scalar;
7#[cfg(target_arch = "x86_64")]
8use hermes_simd_intrinsics::{AmxBf16, Avx512};
9
10impl TiledGemm<Bf16, Bf16, F32> for (Bf16, Bf16, F32) {
11 #[inline]
12 unsafe fn dispatch_tile_matmul(
13 c: *mut F32,
14 c_stride: usize,
15 a: *const Bf16,
16 a_stride: usize,
17 b: *const Bf16,
18 b_stride: usize,
19 ) {
20 #[cfg(target_arch = "x86_64")]
21 {
22 if <Bf16 as AmxSupport>::has_amx() && hermes_simd_intrinsics::AmxSession::is_active() {
23 return <AmxBf16 as TileMatrixMultiply<
24 Bf16,
25 Bf16,
26 F32,
27 AmxBf16,
28 AmxBf16,
29 16,
30 16,
31 32,
32 >>::tile_matmul(c, c_stride, a, a_stride, b, b_stride);
33 }
34 if <Bf16 as Avx512Support>::has_avx512() {
35 return <Avx512 as TileMatrixMultiply<
36 Bf16,
37 Bf16,
38 F32,
39 Avx512,
40 Avx512,
41 16,
42 16,
43 32,
44 >>::tile_matmul(c, c_stride, a, a_stride, b, b_stride);
45 }
46 }
47 <Scalar as TileMatrixMultiply<Bf16, Bf16, F32, Scalar, Scalar, 16, 16, 32>>::tile_matmul(
48 c, c_stride, a, a_stride, b, b_stride,
49 );
50 }
51
52 #[inline]
53 unsafe fn gemm(
54 m: usize,
55 n: usize,
56 k: usize,
57 a: &[Bf16],
58 a_stride: usize,
59 b: &[Bf16],
60 b_stride: usize,
61 c: &mut [F32],
62 c_stride: usize,
63 ) -> Result<(), SimdError> {
64 validate_gemm_sizes(
65 a.len(),
66 b.len(),
67 c.len(),
68 m,
69 n,
70 k,
71 a_stride,
72 b_stride,
73 c_stride,
74 )?;
75
76 #[cfg(target_arch = "x86_64")]
77 {
78 let decision = crate::dispatcher::AdaptiveDispatcher::select_backend(
79 m,
80 n,
81 k,
82 a.as_ptr(),
83 a.len(),
84 b.as_ptr(),
85 b.len(),
86 );
87
88 match decision {
89 crate::dispatcher::DispatchDecision::Amx => {
90 <AmxBf16 as hermes_simd_intrinsics::x86_64::amx::AmxGemm<
91 Bf16,
92 Bf16,
93 F32,
94 >>::amx_gemm(
95 m,
96 n,
97 k,
98 a.as_ptr(),
99 a_stride,
100 b.as_ptr(),
101 b_stride,
102 c.as_mut_ptr(),
103 c_stride,
104 );
105 return Ok(());
106 }
107 crate::dispatcher::DispatchDecision::Avx512 => {
108 tile_loop_generic::<Bf16, Bf16, F32, Avx512, 16, 16, 32>(
109 m,
110 n,
111 k,
112 a.as_ptr(),
113 a_stride,
114 b.as_ptr(),
115 b_stride,
116 c.as_mut_ptr(),
117 c_stride,
118 );
119
120 let amx_m_bound = (m / 16) * 16;
121 let amx_n_bound = (n / 16) * 16;
122 let amx_k_bound = (k / 32) * 32;
123 for r in 0..m {
124 for col in 0..n {
125 if r >= amx_m_bound || col >= amx_n_bound {
126 let mut sum = 0.0f32;
127 for kk in 0..k {
128 sum += a[r * a_stride + kk].to_f32()
129 * b[kk * b_stride + col].to_f32();
130 }
131 c[r * c_stride + col] = F32(c[r * c_stride + col].0 + sum);
132 } else if amx_k_bound < k {
133 let mut sum = 0.0f32;
134 for kk in amx_k_bound..k {
135 sum += a[r * a_stride + kk].to_f32()
136 * b[kk * b_stride + col].to_f32();
137 }
138 c[r * c_stride + col] = F32(c[r * c_stride + col].0 + sum);
139 }
140 }
141 }
142 return Ok(());
143 }
144 crate::dispatcher::DispatchDecision::AvxVnni
145 | crate::dispatcher::DispatchDecision::Scalar => {}
146 }
147 }
148
149 tile_loop_generic::<Bf16, Bf16, F32, Scalar, 16, 16, 32>(
150 m,
151 n,
152 k,
153 a.as_ptr(),
154 a_stride,
155 b.as_ptr(),
156 b_stride,
157 c.as_mut_ptr(),
158 c_stride,
159 );
160
161 let amx_m_bound = (m / 16) * 16;
162 let amx_n_bound = (n / 16) * 16;
163 let amx_k_bound = (k / 32) * 32;
164 for r in 0..m {
165 for col in 0..n {
166 if r >= amx_m_bound || col >= amx_n_bound {
167 let mut sum = 0.0f32;
168 for kk in 0..k {
169 sum += a[r * a_stride + kk].to_f32() * b[kk * b_stride + col].to_f32();
170 }
171 c[r * c_stride + col] = F32(c[r * c_stride + col].0 + sum);
172 } else if amx_k_bound < k {
173 let mut sum = 0.0f32;
174 for kk in amx_k_bound..k {
175 sum += a[r * a_stride + kk].to_f32() * b[kk * b_stride + col].to_f32();
176 }
177 c[r * c_stride + col] = F32(c[r * c_stride + col].0 + sum);
178 }
179 }
180 }
181 Ok(())
182 }
183}