test: code quality testing (#1668)

* fix: improper qualified accesses

* fix: use explicit imports

* fix: actually use explicit imports

* test: add qa testing via Aqua and ExplicitImports

* fix: more ambiguities fixed

* fix: more ambiguities fixed

* fix: more ambiguities fixed

* chore: overlay getindex

* fix: indexing

* fix: more indexing fixes

* feat: all ambiguities have been fixed

* fix: type conversion

* fix: lu factorization

* fix: min compats

* fix: missing getindex for unitrange

* fix: more bug fixes

* fix: getindex ambiguities

* chore: fmt

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* chore: fmt

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix: for 1.10

* fix: use explicit imports

* docs: openssl version mismatch

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

.github/workflows/downgrade.yml

@@ -73,7 +73,7 @@ jobs:
               push!(dev_pks, Pkg.PackageSpec(; path))
           end
           Pkg.develop(dev_pks)
-          Pkg.test(; coverage="user")
+          Pkg.test(; coverage="user", allow_reresolve=false)
         shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
         id: run_tests
         env:

Project.toml

@@ -13,6 +13,7 @@ EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVMOpenMP_jll = "1d63c593-3942-5779-bab2-d838dc0a180e"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -36,13 +37,13 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
-LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
@@ -64,6 +65,7 @@ ReactantOffsetArraysExt = "OffsetArrays"
 ReactantOneHotArraysExt = "OneHotArrays"
 ReactantPythonCallExt = "PythonCall"
 ReactantRandom123Ext = "Random123"
+ReactantSparseArraysExt = "SparseArrays"
 ReactantSpecialFunctionsExt = "SpecialFunctions"
 ReactantStatisticsExt = "Statistics"
 ReactantYaoBlocksExt = "YaoBlocks"
@@ -77,8 +79,8 @@ CUDA = "5.6"
 DLFP8Types = "0.1"
 Downloads = "1.6"
 EnumX = "1"
-Enzyme = "0.13.72"
-EnzymeCore = "0.8.11"
+Enzyme = "0.13.74"
+EnzymeCore = "0.8.13"
 FillArrays = "1.13"
 Float8s = "0.1"
 Functors = "0.5"
@@ -88,14 +90,15 @@ HTTP = "1.10.15"
 KernelAbstractions = "0.9.30"
 LLVM = "9.1"
 LLVMOpenMP_jll = "18.1.7"
+Libdl = "1.10"
 LinearAlgebra = "1.10"
 MPI = "0.20"
 NNlib = "0.9.26"
 OffsetArrays = "1"
 OneHotArrays = "0.2.10"
 OrderedCollections = "1"
 PrecompileTools = "1.2"
-Preferences = "1.4"
+Preferences = "1.4.3"
 PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
@@ -104,6 +107,7 @@ Reactant_jll = "0.0.240"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"
+SparseArrays = "1.10"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
 YaoBlocks = "0.13, 0.14"

docs/Project.toml

@@ -2,6 +2,7 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+OpenSSL_jll = "458c3c95-2e84-50aa-8efc-19380b2a3a95"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -13,3 +14,4 @@ ReactantCore = {path = "../lib/ReactantCore"}
 [compat]
 Documenter = "1.4.1"
 DocumenterVitepress = "0.2"
+OpenSSL_jll = "=3.0.16"

ext/ReactantArrayInterfaceExt.jl

@@ -1,7 +1,7 @@
 module ReactantArrayInterfaceExt
 
 using ArrayInterface: ArrayInterface
-using Reactant: Reactant, RArray, AbstractConcreteNumber, AnyTracedRArray, Ops
+using Reactant: Reactant, RArray, AbstractConcreteNumber, AnyTracedRArray
 
 ArrayInterface.can_setindex(::Type{<:RArray}) = false
 ArrayInterface.fast_scalar_indexing(::Type{<:RArray}) = false

ext/ReactantCUDAExt.jl

@@ -1,23 +1,19 @@
 module ReactantCUDAExt
 
-using CUDA
 using Reactant: Reactant, TracedRArray, AnyConcretePJRTArray, MLIR, TracedRNumber
 using Reactant.Compiler: raising
-using ReactantCore: @trace
+using Reactant.Ops: @opcall
+
+using Adapt: Adapt, adapt
+using CUDA: CUDA, CuDim, DenseCuArray, unsafe_cached_load
+
 using GPUCompiler: GPUCompiler
 using KernelAbstractions: KernelAbstractions
-import KernelAbstractions as KA
 using LLVM: LLVM
-using Libdl
-
-using Reactant.Ops: @opcall
 
-const ReactantKernelAbstractionsExt = Base.get_extension(
-    Reactant, :ReactantKernelAbstractionsExt
-)
-const ReactantBackend = ReactantKernelAbstractionsExt.ReactantBackend
+using PrecompileTools: @setup_workload, @compile_workload
 
-using Adapt
+const KA = KernelAbstractions
 
 Reactant.is_extension_loaded(::Val{:CUDA}) = true
 
@@ -64,9 +60,7 @@ function Base.getindex(RN::CuTracedRNumber{T,A}) where {T,A}
     return @inbounds unsafe_load(RN.ptr, 1, Val(align))
 end
 
-function Base.convert(::Type{T}, RN::CuTracedRNumber) where {T<:Number}
-    return Base.convert(T, Base.getindex(RN))
-end
+Base.convert(::Type{T}, RN::CuTracedRNumber) where {T<:Number} = convert(T, getindex(RN))
 
 for jlop in (
     :(Base.min),
@@ -89,17 +83,15 @@ for jlop in (
     end
 end
 
-@inline Base.ifelse(cond::Bool, a, b::CuTracedRNumber) = Base.ifelse(cond, a, b[])
-@inline Base.ifelse(cond::Bool, a::CuTracedRNumber, b) = Base.ifelse(cond, a[], b)
+@inline Base.ifelse(cond::Bool, a, b::CuTracedRNumber) = ifelse(cond, a, b[])
+@inline Base.ifelse(cond::Bool, a::CuTracedRNumber, b) = ifelse(cond, a[], b)
 @inline Base.ifelse(cond::Bool, a::CuTracedRNumber, b::CuTracedRNumber) =
-    Base.ifelse(cond, a[], b[])
-@inline Base.ifelse(cond::CuTracedRNumber, a, b) = Base.ifelse(cond[], a, b)
-@inline Base.ifelse(cond::CuTracedRNumber, a::CuTracedRNumber, b) =
-    Base.ifelse(cond[], a[], b)
-@inline Base.ifelse(cond::CuTracedRNumber, a, b::CuTracedRNumber) =
-    Base.ifelse(cond[], a, b[])
+    ifelse(cond, a[], b[])
+@inline Base.ifelse(cond::CuTracedRNumber, a, b) = ifelse(cond[], a, b)
+@inline Base.ifelse(cond::CuTracedRNumber, a::CuTracedRNumber, b) = ifelse(cond[], a[], b)
+@inline Base.ifelse(cond::CuTracedRNumber, a, b::CuTracedRNumber) = ifelse(cond[], a, b[])
 @inline Base.ifelse(cond::CuTracedRNumber, a::CuTracedRNumber, b::CuTracedRNumber) =
-    Base.ifelse(cond[], a[], b[])
+    ifelse(cond[], a[], b[])
 
 Base.@constprop :aggressive @inline Base.:^(
     a::CuTracedRNumber{T,A}, b::Integer
@@ -140,7 +132,7 @@ end
                 ),
                 Core.LLVMPtr{UInt8,1},
                 Tuple{Float64},
-                Base.convert(Float64, x),
+                convert(Float64, x),
             ),
         ),
     )
@@ -164,7 +156,7 @@ end
                 ),
                 Core.LLVMPtr{UInt8,1},
                 Tuple{Float32},
-                Base.convert(Float32, x),
+                convert(Float32, x),
             ),
         ),
     )
@@ -181,7 +173,7 @@ Base.@nospecializeinfer function Base.promote_rule(
     @nospecialize(a::Type{<:CuTracedRNumber{T}}),
     @nospecialize(b::Type{<:CuTracedRNumber{T2}})
 ) where {T,T2}
-    return Base.promote_rule(T, T2)
+    return promote_rule(T, T2)
 end
 Base.@nospecializeinfer function Base.promote_rule(
     ::Type{Any}, @nospecialize(b::Type{<:CuTracedRNumber})
@@ -199,7 +191,7 @@ Base.@nospecializeinfer function Base.promote_rule(
     if T == T2
         return T
     else
-        return Base.promote_rule(T, T2)
+        return promote_rule(T, T2)
     end
 end
 Base.@nospecializeinfer function Base.promote_rule(
@@ -208,7 +200,7 @@ Base.@nospecializeinfer function Base.promote_rule(
     if T == T2
         return T
     else
-        return Base.promote_rule(T, T2)
+        return promote_rule(T, T2)
     end
 end
 
@@ -506,9 +498,7 @@ function threads_to_workgroupsize(threads, ndrange)
     end
 end
 
-function ReactantKernelAbstractionsExt.ka_with_reactant(
-    ndrange, workgroupsize, obj, args...
-)
+function Reactant.ka_with_reactant(ndrange, workgroupsize, obj, args...)
     backend = KA.backend(obj)
 
     ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(
@@ -588,7 +578,7 @@ function Adapt.adapt_storage(::ReactantKernelAdaptor, xs::TracedRNumber{T}) wher
     return res
 end
 
-import Reactant.TracedRNumberOverrides.TracedStepRangeLen
+import Reactant.TracedStepRangeLen
 
 function Adapt.adapt_storage(::ReactantKernelAdaptor, r::TracedStepRangeLen)
     return TracedStepRangeLen(
@@ -1481,7 +1471,7 @@ end
 # In Julia v1.11.3 precompiling this module caches bad code:
 # <https://github.com/EnzymeAD/Reactant.jl/issues/614>.
 @static if !Sys.isapple()
-    Reactant.PrecompileTools.@setup_workload begin
+    @setup_workload begin
         Reactant.initialize_dialect()
 
         if Reactant.XLA.REACTANT_XLA_RUNTIME == "PJRT"
@@ -1492,7 +1482,7 @@ end
             error("Unsupported runtime: $(Reactant.XLA.REACTANT_XLA_RUNTIME)")
         end
 
-        Reactant.PrecompileTools.@compile_workload begin
+        @compile_workload begin
             @static if Reactant.precompilation_supported() && VERSION != v"1.11.3"
                 function square_kernel!(x)
                     i = CUDA.threadIdx().x

ext/ReactantKernelAbstractionsExt.jl

@@ -1,14 +1,13 @@
 module ReactantKernelAbstractionsExt
 
-using Reactant
-
-import KernelAbstractions as KA
+using Reactant: Reactant
 
 using Adapt: Adapt
+using KernelAbstractions: KernelAbstractions
 
-## back-end
+const KA = KernelAbstractions
 
-export ReactantBackend
+## back-end
 
 # ToDo: Include XLA client, device and sharding in ReactantBackend struct, to
 # support more complex applications? If so, need to adapt implementation of
@@ -26,7 +25,7 @@ function Base.getproperty(x::ReactantBackend, sym::Symbol)
 end
 
 function KA.allocate(::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
-    return ConcreteRArray{T}(undef, dims)
+    return Reactant.ConcreteRArray{T}(undef, dims)
 end
 
 function KA.zeros(b::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
@@ -103,23 +102,21 @@ end
 
 function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing)
     if Reactant.precompiling()
-        @code_hlo optimize = false tokw(ndrange, workgroupsize, obj, args...)
+        Reactant.@code_hlo optimize = false tokw(ndrange, workgroupsize, obj, args...)
     else
-        @jit tokw(ndrange, workgroupsize, obj, args...)
+        Reactant.@jit tokw(ndrange, workgroupsize, obj, args...)
     end
     return nothing
 end
 
-function ka_with_reactant end # defined in the CUDA extension
-
 Reactant.@reactant_overlay @noinline Base.@nospecializeinfer function (
     obj::KA.Kernel{ReactantBackend}
 )(
     args...; ndrange=nothing, workgroupsize=nothing
 )
     @nospecialize
     return Reactant.call_with_reactant(
-        ka_with_reactant, ndrange, workgroupsize, obj, args...
+        Reactant.ka_with_reactant, ndrange, workgroupsize, obj, args...
     )
 end
 

ext/ReactantNNlibExt/Implementations.jl

@@ -7,7 +7,7 @@ for (jlop, hloop) in (
 end
 
 function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) where {T,N}
-    x = T.(Reactant.materialize_traced_array(x))
+    x = T.(materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = exp.(x .- max_)
     # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581
@@ -22,7 +22,7 @@ function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) whe
 end
 
 function NNlib.logsoftmax!(out::AnyTracedRArray{T}, x::AbstractArray; dims=1) where {T}
-    x = T.(Reactant.materialize_traced_array(x))
+    x = T.(materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = x .- max_
     # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581

ext/ReactantNNlibExt/ReactantNNlibExt.jl

@@ -1,16 +1,13 @@
 module ReactantNNlibExt
 
-using NNlib
-using GPUArraysCore: @allowscalar
 using Reactant:
     Reactant, Ops, TracedRArray, AnyTracedRArray, MLIR, TracedRNumber, @reactant_overlay
-
-using Reactant.TracedUtils:
-    TracedUtils, materialize_traced_array, get_mlir_data, set_mlir_data!
+using Reactant.TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!
 using Reactant.Ops: @opcall
+using ReactantCore: materialize_traced_array, @trace
 
-using ReactantCore: @trace
-using LinearAlgebra: LinearAlgebra, triu
+using NNlib: NNlib, DenseConvDims
+using GPUArraysCore: @allowscalar
 using Statistics: mean
 
 include("Overlay.jl")

ext/ReactantOffsetArraysExt.jl

@@ -1,8 +1,7 @@
 module ReactantOffsetArraysExt
 
-using OffsetArrays
-using OffsetArrays: OffsetArray, OffsetVector
-using Reactant: Reactant, MLIR, Ops, TracedRArray, AbstractConcreteArray
+using OffsetArrays: OffsetArrays, OffsetArray, OffsetVector
+using Reactant: Reactant, MLIR, Ops, TracedRArray, TracedRNumber, AbstractConcreteArray
 
 Base.@nospecializeinfer function Reactant.traced_type_inner(
     @nospecialize(OA::Type{<:OffsetArray}),
@@ -45,22 +44,21 @@ function Base.getindex(
 end
 
 function Base.getindex(
-    x::OffsetVector{Reactant.TracedRNumber{T},Reactant.TracedRArray{T,1}},
-    indices::Base.OneTo{Int},
+    x::OffsetVector{TracedRNumber{T},Reactant.TracedRArray{T,1}}, indices::Base.OneTo{Int}
 ) where {T}
     offset_indices = indices .- x.offsets[1]
     return getindex(parent(x), offset_indices)
 end
 
 parentindex(r::OffsetArrays.IdOffsetRange, i) = i .- r.offset
 function Base.getindex(
-    a::OffsetArray{<:Reactant.TracedRNumber,N}, indices::Vararg{Union{Int,AbstractArray},N}
+    a::OffsetArray{<:TracedRNumber,N}, indices::Vararg{Union{Int,AbstractArray},N}
 ) where {N}
     J = map(parentindex, axes(a), indices)
     return parent(a)[J...]
 end
 
-function Base.getindex(a::OffsetVector{<:Reactant.TracedRNumber}, indices::Int)
+function Base.getindex(a::OffsetVector{<:TracedRNumber}, indices::Int)
     J = parentindex(Base.axes1(a), indices)
     return parent(a)[J]
 end