(
&self,
buf_a: &MappedBuffer,
buf_b: &MappedBuffer,
m: usize,
k: usize,
n: usize,
)
| 942 | } |
| 943 | |
| 944 | fn gpu_gemm( |
| 945 | &self, |
| 946 | buf_a: &MappedBuffer, |
| 947 | buf_b: &MappedBuffer, |
| 948 | m: usize, |
| 949 | k: usize, |
| 950 | n: usize, |
| 951 | ) -> (Vec<f32>, MappedBuffer) { |
| 952 | let buf_c = self.alloc_output(m * n); |
| 953 | // Use small-M kernel (16×64 tile) for M<=16, large kernel (32×64) otherwise |
| 954 | let (entry, wg_m, wg_n) = if m <= 8 { |
| 955 | ("matmul_small", (m as u32).div_ceil(8), (n as u32).div_ceil(64)) |
| 956 | } else { |
| 957 | ("matmul", (m as u32).div_ceil(32), (n as u32).div_ceil(64)) |
| 958 | }; |
| 959 | let result = self.dispatch_compute( |
| 960 | entry, |
| 961 | &[buf_a.buffer, buf_b.buffer, buf_c.buffer], |
| 962 | &buf_c, |
| 963 | m * n, |
| 964 | &[m as u32, n as u32, k as u32, 0], |
| 965 | (wg_m, wg_n, 1), |
| 966 | ); |
| 967 | (result, buf_c) |
| 968 | } |
| 969 | |
| 970 | fn gpu_matmul( |
| 971 | &self, |
no test coverage detected