fix: relax cublaslt types

LuxDL · Oct 10, 2024 · 3c9ff76 · 3c9ff76
1 parent 2d7533c
commit 3c9ff76
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.3.2"
+version = "1.3.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

diff --git a/ext/LuxLibCUDAExt/LuxLibCUDAExt.jl b/ext/LuxLibCUDAExt/LuxLibCUDAExt.jl
@@ -1,7 +1,6 @@
 module LuxLibCUDAExt
 
-# This file only wraps functionality part of CUDA like CUBLAS
-using CUDA: CUDA, CUBLAS, StridedCuMatrix, StridedCuVector, CuPtr, AnyCuMatrix, AnyCuVector
+using CUDA: CUDA, CUBLAS, StridedCuMatrix, StridedCuVector, CuPtr
 using LinearAlgebra: LinearAlgebra, Transpose, Adjoint
 using LuxLib: LuxLib, Optional
 using LuxLib.Utils: ofeltype_array

diff --git a/ext/LuxLibCUDAExt/cublaslt.jl b/ext/LuxLibCUDAExt/cublaslt.jl
@@ -170,16 +170,16 @@ end
 len(x) = length(x)
 len(::Nothing) = nothing
 
-function LuxLib.Impl.cublasLt_fused_dense(act::F, weight::AnyCuMatrix, x::AnyCuMatrix,
-        b::Optional{<:AnyCuVector}, ::False) where {F}
+function LuxLib.Impl.cublasLt_fused_dense(act::F, weight::AbstractMatrix,
+        x::AbstractMatrix, b::Optional{<:AbstractVector}, ::False) where {F}
     z = similar(x, LuxLib.concrete_fba_output_eltype(act, weight, x, b),
         size(weight, 1), size(x, 2))
     LuxLib.cublasLt_fused_dense!(z, act, weight, x, b)
     return z, nothing
 end
 
-function LuxLib.Impl.cublasLt_fused_dense(act::F, weight::AnyCuMatrix, x::AnyCuMatrix,
-        b::Optional{<:AnyCuVector}, ::True) where {F}
+function LuxLib.Impl.cublasLt_fused_dense(act::F, weight::AbstractMatrix,
+        x::AbstractMatrix, b::Optional{<:AbstractVector}, ::True) where {F}
     z = similar(x, LuxLib.concrete_fba_output_eltype(act, weight, x, b),
         size(weight, 1), size(x, 2))
     y = similar(z)
@@ -188,8 +188,8 @@ function LuxLib.Impl.cublasLt_fused_dense(act::F, weight::AnyCuMatrix, x::AnyCuM
 end
 
 function LuxLib.Impl.cublasLt_fused_dense!(
-        z::AbstractMatrix, act::F, weight::AnyCuMatrix, x::AnyCuMatrix,
-        b::Optional{<:AnyCuVector}, y::Optional{<:AbstractMatrix}=nothing) where {F}
+        z::AbstractMatrix, act::F, weight::AbstractMatrix, x::AbstractMatrix,
+        b::Optional{<:AbstractVector}, y::Optional{<:AbstractMatrix}=nothing) where {F}
     if hasmethod(cublaslt_matmul_fused!,
         (typeof(z), typeof(act), typeof(weight), typeof(x), typeof(b), typeof(y)))
         retcode = cublaslt_matmul_fused!(z, act, weight, x, b, y)