fix: missing onehotarrays dispatch for cpu matmul (#1655)

* fix: missing onehotarrays dispatch for cpu matmul

* fix: mooncake lux

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.31.1"
+version = "1.31.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -127,7 +127,7 @@ MLUtils = "0.4.4"
 MPI = "0.20.19"
 MacroTools = "0.5.13"
 Markdown = "1.10"
-Mooncake = "0.4.148, 0.5"
+Mooncake = "0.5"
 NCCL = "0.1.2"
 NNlib = "0.9.27"
 Optimisers = "0.4.6"

ext/MooncakeExt/training.jl

@@ -3,13 +3,7 @@ function Lux.Training.compute_gradients_impl(
 ) where {F}
     config = get_config(ad)
     pullback_cache = prepare_pullback_cache(
-        objective_function,
-        ts.model,
-        ts.parameters,
-        ts.states,
-        data;
-        debug_mode=config.debug_mode,
-        silence_debug_messages=config.silence_debug_messages,
+        objective_function, ts.model, ts.parameters, ts.states, data; config
     )
     # evaluate once to get the correct types
     loss, stₙ, stats = objective_function(ts.model, ts.parameters, ts.states, data)

lib/LuxLib/Project.toml

@@ -1,6 +1,6 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
-version = "1.15.4"
+version = "1.15.3"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]
@@ -36,6 +36,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
+OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
@@ -58,6 +59,7 @@ ForwardDiffExt = "ForwardDiff"
 LoopVectorizationExt = ["LoopVectorization", "Polyester"]
 MKLExt = "MKL"
 OctavianExt = ["Octavian", "LoopVectorization"]
+OneHotArraysExt = ["OneHotArrays"]
 ReactantExt = ["Reactant", "ReactantCore"]
 ReverseDiffExt = "ReverseDiff"
 SLEEFPiratesExt = "SLEEFPirates"
@@ -88,6 +90,7 @@ MLDataDevices = "1.17.1"
 Markdown = "1.10"
 NNlib = "0.9.27"
 Octavian = "0.3.28"
+OneHotArrays = "0.2.5"
 Polyester = "0.7.18"
 Preferences = "1.4.3"
 Random = "1.10"

lib/LuxLib/ext/OneHotArraysExt.jl

@@ -0,0 +1,8 @@
+module OneHotArraysExt
+
+using LuxLib: Utils
+using OneHotArrays: OneHotLike
+
+Utils.force_3arg_mul!_dispatch(::AbstractMatrix, ::AbstractMatrix, ::OneHotLike) = true
+
+end

lib/LuxLib/src/impl/matmul.jl

@@ -109,6 +109,10 @@ function matmuladd!(
 end
 
 function matmul!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)
+    if Utils.force_3arg_mul!_dispatch(C, A, B)
+        mul!(C, A, B)
+        return nothing
+    end
     matmul!(C, internal_operation_mode((C, A, B)), A, B)
     return nothing
 end

lib/LuxLib/src/utils.jl

@@ -396,4 +396,8 @@ function maybe_reshape(x::AbstractArray{T,N}, dims::Dims{N}) where {T,N}
 end
 maybe_reshape(x::AbstractArray{T}, dims::Dims) where {T} = reshape(x, dims)
 
+force_3arg_mul!_dispatch(::AbstractMatrix, ::AbstractMatrix, ::AbstractMatrix) = false
+
+CRC.@non_differentiable force_3arg_mul!_dispatch(::Any...)
+
 end

lib/LuxLib/test/Project.toml

@@ -22,6 +22,7 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
+OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

lib/LuxLib/test/others/onehotarray_tests.jl

@@ -0,0 +1,11 @@
+using OneHotArrays, LuxLib, Test
+
+@testset "Specialized OneHotArrays Dispatch" begin
+    x = onehotbatch("aabc", "abcdefghijklmnopqrstuv")
+    weight = reshape(collect(Float32, 1:(1024 * 22)), 1024, 22)
+
+    @test fused_dense_bias_activation(identity, weight, x, nothing) ≈
+        fused_dense_bias_activation(identity, weight, Array(x), nothing)
+
+    @test LuxLib.Utils.force_3arg_mul!_dispatch(weight, weight, x)
+end