(&self, a: &Tensor, b: &Tensor)
| 280 | } |
| 281 | |
| 282 | fn tensor_matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> { |
| 283 | let a = if a.dtype() == DType::F32 && a.is_contiguous() { a.clone() } |
| 284 | else { a.to_dtype(DType::F32)?.contiguous()? }; |
| 285 | let b = if b.dtype() == DType::F32 && b.is_contiguous() { b.clone() } |
| 286 | else { b.to_dtype(DType::F32)?.contiguous()? }; |
| 287 | let ad = a.dims(); let bd = b.dims(); |
| 288 | let ra = ad.len(); let rb = bd.len(); |
| 289 | let m = ad[ra-2]; let k = ad[ra-1]; let n = bd[rb-1]; |
| 290 | let ab: usize = ad[..ra-2].iter().product(); |
| 291 | let bb: usize = bd[..rb-2].iter().product(); |
| 292 | let batch = ab.max(bb); |
| 293 | |
| 294 | if batch <= 1 { |
| 295 | let ap = self.get_or_upload(&a)?; |
| 296 | let bp = self.get_or_upload(&b)?; |
| 297 | let c = self.rocblas_gemm(ap, bp, m, k, n)?; |
| 298 | let data = c.download_on_stream(self.stream).map_err(candle_core::Error::Msg)?; |
| 299 | let mut s = ad[..ra-2].to_vec(); s.push(m); s.push(n); |
| 300 | return Tensor::from_vec(data, s.as_slice(), &Device::Cpu); |
| 301 | } |
| 302 | |
| 303 | let ad_flat = Self::to_f32_vec(&a)?; |
| 304 | let bd_flat = Self::to_f32_vec(&b)?; |
| 305 | let (mk, kn, mn) = (m*k, k*n, m*n); |
| 306 | // Upload entire A and B arrays at once — one H2D transfer per operand |
| 307 | let a_gpu = self.gpu_upload(&ad_flat)?; |
| 308 | let b_gpu = self.gpu_upload(&bd_flat)?; |
| 309 | let c_gpu = self.gpu_alloc(batch * mn)?; |
| 310 | let alpha: f32 = 1.0; |
| 311 | let beta: f32 = 0.0; |
| 312 | // Use strided batched GEMM when both operands have uniform strides |
| 313 | if ab == batch && bb == batch { |
| 314 | let err = unsafe { |
| 315 | (self.ffi.rocblas_sgemm_strided_batched)( |
| 316 | self.blas_handle, |
| 317 | ROCBLAS_OPERATION_NONE, ROCBLAS_OPERATION_NONE, |
| 318 | n as i32, m as i32, k as i32, |
| 319 | &alpha, |
| 320 | b_gpu.as_ptr(), n as i32, kn as i64, |
| 321 | a_gpu.as_ptr(), k as i32, mk as i64, |
| 322 | &beta, |
| 323 | c_gpu.as_mut_ptr(), n as i32, mn as i64, |
| 324 | batch as i32, |
| 325 | ) |
| 326 | }; |
| 327 | if err != 0 { return Err(candle_core::Error::Msg(format!("rocblas_sgemm_strided_batched: {err}"))); } |
| 328 | } else { |
| 329 | // Fallback: loop for broadcast batches (one operand has batch=1) |
| 330 | for i in 0..batch { |
| 331 | let ao = if ab==1 {0} else {i*mk}; |
| 332 | let bo = if bb==1 {0} else {i*kn}; |
| 333 | let co = i * mn; |
| 334 | let err = unsafe { |
| 335 | (self.ffi.rocblas_sgemm)( |
| 336 | self.blas_handle, |
| 337 | ROCBLAS_OPERATION_NONE, ROCBLAS_OPERATION_NONE, |
| 338 | n as i32, m as i32, k as i32, |
| 339 | &alpha, |
no test coverage detected