Replace the fallback `mul!` in `batchedmul` with `Octavian.matmul!`?
- Dominant language
- Julia
- Stars
- 258
- Forks
- 138
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 3
Description
Currently the `batchedmul` will [fallback to calling `mul!`](https://github.com/FluxML/NNlib.jl/blob/0b64dc11e6ba47707c43cf668663e48a615c85bb/src/batched/batchedmul.jl#L275) if the element type doesn't have BLAS support, and the fallback `mul!` is quite slow. For example, `batchedmul(randn(Float16, 1024, 1024, 32), randn(Float16, 1024, 1024, 32))` on my computer never terminated within 5 minutes, while replace the line with `Octavian.matmul!` and re-run give:
```julia
julia> a = randn(Float16, 1024, 1024, 32);
julia> b = randn(Float16, 1024, 1024, 32);
julia> @btime batched_mul($a, $b);
260.851 ms (2 allocations: 64.00 MiB)
julia> a32 = randn(Float32, 1024, 1024, 32);
julia> b32 = randn(Float32, 1024, 1024, 32);
julia> @btime batched_mul($a32, $b32);
128.949 ms (38 allocations: 128.00 MiB)
```
Which is not optimal, but IMPO good enough compare to the current status.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.