Remove ROCKernels dep

exanauts · Nov 18, 2022 · 8f4e5e9 · 8f4e5e9
1 parent 5d57876
commit 8f4e5e9
Show file tree

Hide file tree

Showing 16 changed files with 36 additions and 37 deletions.
diff --git a/Project.toml b/Project.toml
@@ -18,7 +18,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-AMDGPU = "0.4"
+AMDGPU = "=0.4.2"
 CUDA = "3.4"
 ExaTron = "2"
 FileIO = "1.14"

diff --git a/src/algorithms/admm_two_level.jl b/src/algorithms/admm_two_level.jl
@@ -1,5 +1,5 @@
 function admm_two_level(
-    env::AdmmEnv, mod::AbstractOPFModel, device::Union{Nothing,KA.GPU}=nothing
+    env::AdmmEnv, mod::AbstractOPFModel, device=nothing
 )
     par = env.params
     info = mod.info

diff --git a/src/models/acopf/acopf_admm_prepoststep_ka.jl b/src/models/acopf/acopf_admm_prepoststep_ka.jl
@@ -1,7 +1,7 @@
 function admm_outer_prestep(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     sol, info = mod.solution, mod.info
     info.norm_z_prev = norm(sol.z_curr)
@@ -11,7 +11,7 @@ end
 function admm_inner_prestep(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     sol = mod.solution
     ev = copy_data_kernel_ka(device, 64, mod.nvar)(
@@ -26,7 +26,7 @@ end
 function admm_poststep(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     data, sol, info, grid_data = env.data, mod.solution, mod.info, mod.grid_data
 

diff --git a/src/models/acopf/acopf_admm_rolling_ka.jl b/src/models/acopf/acopf_admm_rolling_ka.jl
@@ -19,7 +19,7 @@ end
 function admm_restart_rolling(
     env::AdmmEnv,
     mod::ModelAcopf,
-    device::KA.GPU,
+    device,
     start_period=1, end_period=6, result_file="warm-start")
 
     @assert env.load_specified == true

diff --git a/src/models/acopf/acopf_admm_update_l_ka.jl b/src/models/acopf/acopf_admm_update_l_ka.jl
@@ -16,7 +16,7 @@ end
 function admm_update_l(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     par, sol, info = env.params, mod.solution, mod.info
     ev = update_l_kernel_ka(device,64,mod.nvar)(

diff --git a/src/models/acopf/acopf_admm_update_lz_ka.jl b/src/models/acopf/acopf_admm_update_lz_ka.jl
@@ -13,7 +13,7 @@ end
 function admm_update_lz(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     par, sol, info = env.params, mod.solution, mod.info
     ev = update_lz_kernel_ka(device,64,mod.nvar)(

diff --git a/src/models/acopf/acopf_admm_update_penalty_ka.jl b/src/models/acopf/acopf_admm_update_penalty_ka.jl
@@ -15,7 +15,7 @@ end
 function acopf_admm_update_penalty(
     env::AdmmEnv,
     mod::Model,
-    device::KA.GPU
+    device
 )
 #=
     par, sol = env.params, mod.solution

diff --git a/src/models/acopf/acopf_admm_update_residual_ka.jl b/src/models/acopf/acopf_admm_update_residual_ka.jl
@@ -13,7 +13,7 @@ end
 function admm_update_residual(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     sol, info = mod.solution, mod.info
 

diff --git a/src/models/acopf/acopf_admm_update_x_ka.jl b/src/models/acopf/acopf_admm_update_x_ka.jl
@@ -2,7 +2,7 @@ function acopf_admm_update_x_gen(
     env::AdmmEnv,
     mod::AbstractOPFModel,
     gen_solution::EmptyGeneratorSolution,
-    device::KA.GPU
+    device
 )
     sol, info, data = mod.solution, mod.info, mod.grid_data
     generator_kernel_two_level(mod, data.baseMVA, sol.u_curr, sol.v_curr, sol.z_curr, sol.l_curr, sol.rho, device)
@@ -13,7 +13,7 @@ end
 function acopf_admm_update_x_line(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     par, sol, info, data = env.params, mod.solution, mod.info, mod.grid_data
     if env.use_linelimit
@@ -45,7 +45,7 @@ end
 function admm_update_x(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     acopf_admm_update_x_gen(env, mod, mod.gen_solution, device)
     acopf_admm_update_x_line(env, mod, device)

diff --git a/src/models/acopf/acopf_admm_update_xbar_ka.jl b/src/models/acopf/acopf_admm_update_xbar_ka.jl
@@ -1,7 +1,7 @@
 function admm_update_xbar(
     env::AdmmEnv,
     mod::AbstractOPFModel,
-    device::KA.GPU
+    device
 )
     sol, info, data = mod.solution, mod.info, mod.grid_data
     # nblk_bus = div(data.nbus, 32, RoundUp)

diff --git a/src/models/acopf/acopf_generator_kernel_ka.jl b/src/models/acopf/acopf_generator_kernel_ka.jl
@@ -26,7 +26,7 @@ end
 function generator_kernel_two_level(
     model::ModelAcopf,
     baseMVA::Float64, u, xbar,
-    zu, lu, rho_u, device::KA.GPU
+    zu, lu, rho_u, device
 )
     nblk = div(model.grid_data.ngen, 32, RoundUp)
     ev = generator_kernel_two_level_ka(device,32,nblk*32)(
@@ -37,4 +37,4 @@ function generator_kernel_two_level(
     )
     wait(ev)
     return 0.0
-end
+end
diff --git a/src/models/acopf/acopf_init_solution_ka.jl b/src/models/acopf/acopf_init_solution_ka.jl
@@ -53,7 +53,7 @@ end
 function init_solution!(
     model::AbstractOPFModel,
     sol::Solution,
-    rho_pq::Float64, rho_va::Float64, device::KA.GPU
+    rho_pq::Float64, rho_va::Float64, device
 )
     fill!(sol, 0.0)
     sol.rho .= rho_pq

diff --git a/src/models/acopf/acopf_set_linelimit_ka.jl b/src/models/acopf/acopf_set_linelimit_ka.jl
@@ -14,7 +14,7 @@ function acopf_set_linelimit(
     env::AdmmEnv,
     mod::AbstractOPFModel,
     info::IterationInformation,
-    device::KA.GPU
+    device
 )
     ev = set_rateA_kernel_ka(device,64,mod.nline)(mod.nline, mod.membuf, mod.rateA)
     wait(ev)

diff --git a/src/utils/environment.jl b/src/utils/environment.jl
@@ -91,7 +91,7 @@ mutable struct AdmmEnv{T,TD,TI,TM} <: AbstractAdmmEnv{T,TD,TI,TM}
     tight_factor::T
     horizon_length::Int
     use_gpu::Bool
-    ka_device::Union{Nothing,KA.GPU}
+    ka_device
     use_linelimit::Bool
     use_mpi::Bool
     use_projection::Bool

diff --git a/src/utils/opfdata.jl b/src/utils/opfdata.jl
@@ -578,7 +578,7 @@ function get_generator_data(data::OPFData, device::Nothing=nothing; use_gpu=fals
   return pgmin,pgmax,qgmin,qgmax,c2,c1,c0
 end
 
-function get_generator_data(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_generator_data(data::OPFData, device; use_gpu=false)
   ngen = length(data.generators)
 
     pgmin = KAArray{Float64}(ngen, device)
@@ -651,7 +651,7 @@ function get_bus_data(data::OPFData, device::Nothing=nothing; use_gpu=false)
   end
 end
 
-function get_bus_data(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_bus_data(data::OPFData, device; use_gpu=false)
   nbus = length(data.buses)
 
   FrIdx = Int[l for b=1:nbus for l in data.FromLines[b]]
@@ -755,7 +755,7 @@ function get_branch_data(data::OPFData, device::Nothing=nothing; use_gpu::Bool=f
   end
 end
 
-function get_branch_data(data::OPFData, device::KA.GPU; use_gpu::Bool=false, tight_factor::Float64=1.0)
+function get_branch_data(data::OPFData, device; use_gpu::Bool=false, tight_factor::Float64=1.0)
   buses = data.buses
   lines = data.lines
   BusIdx = data.BusIdx
@@ -829,7 +829,7 @@ function get_branch_bus_index(data::OPFData, device::Nothing=nothing; use_gpu=fa
   end
 end
 
-function get_branch_bus_index(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_branch_bus_index(data::OPFData, device; use_gpu=false)
   lines = data.lines
   BusIdx = data.BusIdx
   nline = length(lines)
@@ -866,7 +866,7 @@ function get_generator_bus_data(data::OPFData, device::Nothing=nothing; use_gpu=
   return vgmin, vgmax, vm_setpoint
 end
 
-function get_generator_bus_data(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_generator_bus_data(data::OPFData, device; use_gpu=false)
   ngen = length(data.generators)
 
   vgmin = KAArray{Float64}(ngen, device)
@@ -906,7 +906,7 @@ function get_generator_primary_control(data::OPFData, device::Nothing=nothing; d
   return alpha_g, pg_setpoint
 end
 
-function get_generator_primary_control(data::OPFData, device::KA.GPU; droop::Float64=0.04, use_gpu=false)
+function get_generator_primary_control(data::OPFData, device; droop::Float64=0.04, use_gpu=false)
   ngen = length(data.generators)
 
   alpha_g = KAArray{Float64}(ngen, device)
@@ -955,7 +955,7 @@ function get_storage_data(data::OPFData, device::Nothing=nothing; use_gpu=false)
   end
 end
 
-function get_storage_data(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_storage_data(data::OPFData, device; use_gpu=false)
   nstorage = length(data.storages)
 
   chg_min = Float64[data.storages[s].chg_min for s=1:nstorage]
@@ -1004,7 +1004,7 @@ function get_bus_storage_index(data::OPFData, device::Nothing=nothing; use_gpu=f
   end
 end
 
-function get_bus_storage_index(data::OPFData, device::KA.GPU; use_gpu=false)
+function get_bus_storage_index(data::OPFData, device; use_gpu=false)
   nbus = length(data.buses)
 
   StorageIdx = Int[s for b=1:nbus for s in data.BusStorages[b]]

diff --git a/test/algorithms/acopf_update_ka.jl b/test/algorithms/acopf_update_ka.jl
@@ -2,16 +2,15 @@ using CUDA
 using AMDGPU
 using KernelAbstractions
 KA = KernelAbstractions
-devices = Vector{KA.Device}()
-push!(devices, KA.CPU())
+devices = []
 if CUDA.has_cuda_gpu() || AMDGPU.has_rocm_gpu()
-    if CUDA.has_cuda_gpu()
-        using CUDAKernels
-        function ExaAdmm.KAArray{T}(n::Int, device::CUDADevice) where {T}
-            return CuArray{T}(undef, n)
-        end
-        push!(devices, CUDADevice())
-    end
+   if CUDA.has_cuda_gpu()
+       using CUDAKernels
+       function ExaAdmm.KAArray{T}(n::Int, device::CUDADevice) where {T}
+           return CuArray{T}(undef, n)
+       end
+       push!(devices, CUDADevice())
+   end
     if AMDGPU.has_rocm_gpu()
         using ROCKernels
         function ExaAdmm.KAArray{T}(n::Int, device::ROCDevice) where {T}
@@ -32,7 +31,7 @@ end
     if isa(device, KA.CPU)
         TD = Array{Float64,1}; TI = Array{Int,1}; TM = Array{Float64,2}
         env = ExaAdmm.AdmmEnv{T,TD,TI,TM}(case, rho_pq, rho_va; use_gpu=false, ka_device=device, verbose=verbose)
-    else isa(device, KA.GPU)
+    else isa(device, ROCDevice)
         if CUDA.has_cuda_gpu()
             TD = CuArray{Float64,1}; TI = CuArray{Int,1}; TM = CuArray{Float64,2}
         elseif AMDGPU.has_rocm_gpu()