fix: missing onehotarrays dispatch for cpu matmul

Author: avik-pal

lib/LuxLib/Project.toml

@@ -1,6 +1,6 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
-version = "1.15.4"
+version = "1.15.5"
 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,15 @@
+using OneHotArrays, LuxLib, Test, BenchmarkTools
+
+@testset "Specialized OneHotArrays Dispatch" begin
+    x = onehotbatch("aabc", "abcdefghijklmnopqrstuv")
+    weight = reshape(collect(Float32, 1:(1024 * 22)), 1024, 22)
+
+    dense_res_time = @belapsed fused_dense_bias_activation(
+        identity, weight, Array(x), nothing
+    )
+    onehot_res_time = @belapsed fused_dense_bias_activation(identity, weight, x, nothing)
+
+    @test onehot_res_time < dense_res_time / 5
+    @test fused_dense_bias_activation(identity, weight, x, nothing) ≈
+        fused_dense_bias_activation(identity, weight, Array(x), nothing)
+end