FluxML / FluxML/NNlib.jl

Replace the fallback `mul!` in `batchedmul` with `Octavian.matmul!`?

Open
#446 11 comments 0 reactions 0 assignees View on GitHub
performance
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.