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

[chengchingwen]

Currently the batchedmul will fallback to calling mul! 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> 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.

[CarloLucibello]

If we do this, should we rely on Octvian for all float types? Would it be faster?

[chengchingwen]

I'm not sure whether the multi-thread strategy we used with BLAS gemm would work well with the Octavian's multi-thread strategy, add we probably want that to live in Octavian (JuliaLinearAlgebra/Octavian.jl#40).

[ToucheSir]

How are load times and TTFX for Octavian these days? As I recall that was a huge blocker for any package seriously depending on it.

[chengchingwen]

On julia 1.8.3

julia> @time using Octavian                             
 14.340930 seconds (10.25 M allocations: 900.762 MiB, 7.47% gc time, 39.17% compilation time: 50% of which was recompilation)

julia> @time using Octavian
  0.000147 seconds (122 allocations: 10.430 KiB)

julia> @time matmul(randn(3,3), randn(3,3));
  0.000013 seconds (3 allocations: 384 bytes)

julia> @time matmul(randn(3,3), randn(3,3));
  0.000006 seconds (3 allocations: 384 bytes)

julia> @time matmul(randn(Float16, 3,3), randn(Float16, 3,3));
  0.000005 seconds (3 allocations: 240 bytes)

julia> @time matmul(randn(Float16, 3,3), randn(Float16, 3,3));
  0.000007 seconds (3 allocations: 240 bytes)

[ToucheSir]

That's more time than Flux itself 😢. We could look into adding it as a package extension once 1.9 releases.

[chengchingwen]

Would load time and TTFX be a huge blocker for packages like Flux which the functions would be run in multiple times and the average running time is far more larger than seconds?

[ToucheSir]

Yes. It makes development an absolute pain for one, but also more users are running script-based workflows than the average Julia userbase. TTFX is bad enough for Flux models already, so ideally we wouldn't add more on. I am confused by why Octavian's import time is so slow, has anyone looked into that?

[chengchingwen]

I am confused by why Octavian's import time is so slow, has anyone looked into that?

I didn't see any issues about it, but:

julia> @time_imports using Octavian
      0.3 ms  Requires           
      3.3 ms  ArrayInterfaceCore
      0.1 ms  IfElse     
     41.2 ms  Static     
      0.7 ms  Compat      
      5.0 ms  ArrayInterface
      0.3 ms  Adapt             
     81.0 ms  OffsetArrays    
      0.7 ms  ArrayInterfaceOffsetArrays
      2.7 ms  StaticArraysCore
    585.2 ms  StaticArrays                
      0.4 ms  ArrayInterfaceStaticArraysCore         
      1.3 ms  ArrayInterfaceStaticArrays
      0.1 ms  SIMDTypes 
      1.7 ms  ManualMemory        
      3.9 ms  LayoutPointers
      1.1 ms  CpuId          
    192.4 ms  CPUSummary 91.92% compilation time
      0.1 ms  BitTwiddlingConvenienceFunctions                    
    124.2 ms  HostCPUFeatures 61.10% compilation time
    482.6 ms  VectorizationBase
      4.3 ms  SLEEFPirates
      0.2 ms  UnPack
     50.5 ms  ThreadingUtilities 77.56% compilation time
     18.6 ms  PolyesterWeave 68.95% compilation time
      3.1 ms  DocStringExtensions 53.93% compilation time
      4.4 ms  CloseOpenIntervals
      0.2 ms  SnoopPrecompile
     74.4 ms  ChainRulesCore
      5.8 ms  IrrationalConstants
      0.9 ms  DiffRules
      2.4 ms  DiffResults
     23.5 ms  Preferences
      0.2 ms  OpenLibm_jll
      0.3 ms  NaNMath
     13.4 ms  ChangesOfVariables
      1.3 ms  InverseFunctions
      0.8 ms  LogExpFunctions
      0.3 ms  JLLWrappers
      0.3 ms  CompilerSupportLibraries_jll
      2.9 ms  OpenSpecFun_jll 83.04% compilation time
     16.6 ms  SpecialFunctions
     78.0 ms  MacroTools
      0.3 ms  CommonSubexpressions
    148.8 ms  ForwardDiff
      3.1 ms  SIMDDualNumbers
    202.3 ms  LoopVectorization
   5624.9 ms  Octavian 83.46% compilation time (56% recompilation)

There is a large recompilation on Octavian itself

[mcabbott]

Now that we have package extensions (as in e.g. #445) we could fairly easily have optional faster implementations, which are enabled by loading e.g. using Octavian yourself. (Since Float16 is so slow, the base implementation could even just warn you that you should do this.)

Is Octavian obviously the right level? It does clever things to multi-thread etc. a single large matmul. But in the batched case, much of that would ideally apply to the whole batched operation instead. A much simpler set of 4 loops with @tturbo might be competitive.

And an extension enabled by using LoopVectorization would be nice for other things, like softmax (can do better than #450) and perhaps conv too (see FluxML/NNlibCPU.jl#1).

[chengchingwen]

Is Octavian obviously the right level? It does clever things to multi-thread etc. a single large matmul. But in the batched case, much of that would ideally apply to the whole batched operation instead. A much simpler set of 4 loops with @tturbo might be competitive.

It would be better to add batch support directly into Octavian as JuliaLinearAlgebra/Octavian.jl#40.

[mcabbott]

Ah, as linked above. That would be even better, if someone gets around to writing it.

(Also better, for Float32, would probably be to call MKL's batched routine instead of the loop with mul!. But nobody has hooked that up so far.)