Migrate to KA 0.9

exanauts · Apr 19, 2023 · e62f009 · e62f009
1 parent c92a909
commit e62f009
Show file tree

Hide file tree

Showing 16 changed files with 54 additions and 71 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ExaAdmm"
 uuid = "4d6a948c-1075-4240-a564-361a5d4e22a2"
 authors = ["Youngdae Kim <[email protected]>", "Kibaek Kim <[email protected]>", "Weiqi Zhang <[email protected]>", "François Pacaud <[email protected]>", "Michel Schanen <[email protected]>"]
-version = "0.3.2"
+version = "0.4.0"
 
 [deps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
@@ -19,9 +19,9 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 AMDGPU = "0.4"
-CUDA = "3.4"
-ExaTron = "2"
+CUDA = "4"
+ExaTron = "3"
 FileIO = "1.14"
-KernelAbstractions = "0.8"
-MPI = "0.19"
+KernelAbstractions = "0.9"
+MPI = "0.20"
 julia = "1.8"
diff --git a/README.md b/README.md
@@ -37,9 +37,8 @@ and the same example on an AMD GPU:
 ```julia
 using ExaAdmm
 using AMDGPU
-using ROCKernels
 
-ExaAdmm.KAArray{T}(n::Int, ::ROCDevice) where {T} = ROCArray{T}(undef, n)
+ExaAdmm.KAArray{T}(n::Int, ::ROCBackend) where {T} = ROCArray{T}(undef, n)
 
 env, mod = solve_acopf(
     "case1354pegase.m";
@@ -50,7 +49,7 @@ env, mod = solve_acopf(
     scale=1e-4,
     tight_factor=0.99,
     use_gpu=true,
-    ka_device = ROCDevice(),
+    ka_device = ROCBackend(),
     verbose=1
 )
 ```

diff --git a/src/models/acopf/acopf_admm_prepoststep_ka.jl b/src/models/acopf/acopf_admm_prepoststep_ka.jl
@@ -14,11 +14,10 @@ function admm_inner_prestep(
     device
 )
     sol = mod.solution
-    ev = copy_data_kernel_ka(device, 64, mod.nvar)(
-        mod.nvar, sol.z_prev, sol.z_curr,
-        dependencies=Event(device)
+    copy_data_kernel_ka(device, 64, mod.nvar)(
+        mod.nvar, sol.z_prev, sol.z_curr
     )
-    wait(ev)
+    KA.synchronize(device)
 
     return
 end

diff --git a/src/models/acopf/acopf_admm_rolling_ka.jl b/src/models/acopf/acopf_admm_rolling_ka.jl
@@ -70,13 +70,12 @@ function admm_restart_rolling(
         #@printf(io, "Line violations (RateA) . . . . . . . . . %.6e\n", mod.solution.max_line_viol_rateA)
         @printf(io, "Time (secs) . . . . . . . . . . . . . . . %5.3f\n", mod.info.time_overall + mod.info.time_projection)
 
-        ev = update_real_power_current_bounds_ka(device,64,ngen)(
+        update_real_power_current_bounds_ka(device,64,ngen)(
             mod.grid_data.ngen, mod.gen_start,
             mod.pgmin_curr, mod.pgmax_curr, mod.grid_data.pgmin, mod.grid_data.pgmax,
-            mod.grid_data.ramp_rate, mod.solution.u_curr,
-            dependencies=Event(device)
+            mod.grid_data.ramp_rate, mod.solution.u_curr
         )
-        wait(ev)
+        KA.synchronize(device)
 
         flush(io)
     end

diff --git a/src/models/acopf/acopf_admm_update_l_ka.jl b/src/models/acopf/acopf_admm_update_l_ka.jl
@@ -22,7 +22,7 @@ function admm_update_l(
     ev = update_l_kernel_ka(device,64,mod.nvar)(
         mod.nvar, sol.l_curr, sol.z_curr, sol.lz, par.beta
     )
-    wait(ev)
+    KA.synchronize(device)
     info.time_l_update += 0.0
     return
 end
diff --git a/src/models/acopf/acopf_admm_update_lz_ka.jl b/src/models/acopf/acopf_admm_update_lz_ka.jl
@@ -16,11 +16,10 @@ function admm_update_lz(
     device
 )
     par, sol, info = env.params, mod.solution, mod.info
-    ev = update_lz_kernel_ka(device,64,mod.nvar)(
-        mod.nvar, par.MAX_MULTIPLIER, sol.z_curr, sol.lz, par.beta,
-        dependencies=Event(device)
+    update_lz_kernel_ka(device,64,mod.nvar)(
+        mod.nvar, par.MAX_MULTIPLIER, sol.z_curr, sol.lz, par.beta
     )
-    wait(ev)
+    KA.synchronize(device)
     info.time_lz_update += 0.0
     return
 end
diff --git a/src/models/acopf/acopf_admm_update_residual_ka.jl b/src/models/acopf/acopf_admm_update_residual_ka.jl
@@ -17,21 +17,18 @@ function admm_update_residual(
 )
     sol, info = mod.solution, mod.info
 
-    ev = compute_primal_residual_kernel_ka(device,64,mod.nvar)(
-        mod.nvar, sol.rp, sol.u_curr, sol.v_curr, sol.z_curr,
-        dependencies=Event(device)
+    compute_primal_residual_kernel_ka(device,64,mod.nvar)(
+        mod.nvar, sol.rp, sol.u_curr, sol.v_curr, sol.z_curr
     )
-    wait(ev)
-    ev = vector_difference_ka(device,64,mod.nvar)(
-        mod.nvar, sol.rd, sol.z_curr, sol.z_prev,
-        dependencies=Event(device)
+    KA.synchronize(device)
+    vector_difference_ka(device,64,mod.nvar)(
+        mod.nvar, sol.rd, sol.z_curr, sol.z_prev
     )
-    wait(ev)
-    ev = vector_difference_ka(device,64,mod.nvar)(
-        mod.nvar, sol.Ax_plus_By, sol.rp, sol.z_curr,
-        dependencies=Event(device)
+    KA.synchronize(device)
+    vector_difference_ka(device,64,mod.nvar)(
+        mod.nvar, sol.Ax_plus_By, sol.rp, sol.z_curr
     )
-    wait(ev)
+    KA.synchronize(device)
 
     info.primres = norm(sol.rp, device)
     info.dualres = norm(sol.rd, device)

diff --git a/src/models/acopf/acopf_admm_update_x_ka.jl b/src/models/acopf/acopf_admm_update_x_ka.jl
@@ -17,25 +17,23 @@ function acopf_admm_update_x_line(
 )
     par, sol, info, data = env.params, mod.solution, mod.info, mod.grid_data
     if env.use_linelimit
-        ev = auglag_linelimit_two_level_alternative_ka(device, 32, data.nline*32)(
+        auglag_linelimit_two_level_alternative_ka(device, 32, data.nline*32)(
             Val(mod.n), data.nline, mod.line_start,
             info.inner, par.max_auglag, par.mu_max, par.scale,
             sol.u_curr, sol.v_curr, sol.z_curr, sol.l_curr, sol.rho,
             par.shift_lines, mod.membuf, data.YffR, data.YffI, data.YftR, data.YftI,
             data.YttR, data.YttI, data.YtfR, data.YtfI,
-            data.FrVmBound, data.ToVmBound, data.FrVaBound, data.ToVaBound,
-            dependencies=Event(device)
+            data.FrVmBound, data.ToVmBound, data.FrVaBound, data.ToVaBound
         )
-        wait(ev)
+        KA.synchronize(device)
     else
-        ev = polar_kernel_two_level_alternative_ka(device, 32, data.nline)(
+        polar_kernel_two_level_alternative_ka(device, 32, data.nline)(
             mod.n, data.nline, mod.line_start, par.scale,
             sol.u_curr, sol.v_curr, sol.z_curr, sol.l_curr, sol.rho,
             par.shift_lines, mod.membuf, data.YffR, data.YffI, data.YftR, data.YftI,
-            data.YttR, data.YttI, data.YtfR, data.YtfI, data.FrVmBound, data.ToVmBound,
-            dependencies=Event(device)
+            data.YttR, data.YttI, data.YtfR, data.YtfI, data.FrVmBound, data.ToVmBound
         )
-        wait(ev)
+        KA.synchronize(device)
     end
     info.time_x_update += 0.0
     info.user.time_branches += 0.0

diff --git a/src/models/acopf/acopf_admm_update_xbar_ka.jl b/src/models/acopf/acopf_admm_update_xbar_ka.jl
@@ -6,14 +6,13 @@ function admm_update_xbar(
     sol, info, data = mod.solution, mod.info, mod.grid_data
     # nblk_bus = div(data.nbus, 32, RoundUp)
 
-    ev = bus_kernel_two_level_alternative_ka(device,32,data.nbus)(
+    bus_kernel_two_level_alternative_ka(device,32,data.nbus)(
         data.baseMVA, data.nbus, mod.gen_start, mod.line_start,
         data.FrStart, data.FrIdx, data.ToStart, data.ToIdx, data.GenStart,
         data.GenIdx, data.Pd, data.Qd, sol.u_curr, sol.v_curr,
-        sol.z_curr, sol.l_curr, sol.rho, data.YshR, data.YshI,
-        dependencies=Event(device)
+        sol.z_curr, sol.l_curr, sol.rho, data.YshR, data.YshI
     )
-    wait(ev)
+    KA.synchronize(device)
     info.time_xbar_update += 0.0
     info.user.time_buses += 0.0
     return

diff --git a/src/models/acopf/acopf_admm_update_z_ka.jl b/src/models/acopf/acopf_admm_update_z_ka.jl
@@ -17,12 +17,11 @@ function admm_update_z(
 )
     par, sol, info = env.params, mod.solution, mod.info
 
-    ev = update_zv_kernel_ka(device,64,mod.nvar)(
+    update_zv_kernel_ka(device,64,mod.nvar)(
         mod.nvar, sol.u_curr, sol.v_curr, sol.z_curr,
-        sol.l_curr, sol.rho, sol.lz, par.beta,
-        dependencies=Event(device)
+        sol.l_curr, sol.rho, sol.lz, par.beta
     )
-    wait(ev)
+    KA.synchronize(device)
     info.time_z_update += 0.0
     return
 end
diff --git a/src/models/acopf/acopf_generator_kernel_ka.jl b/src/models/acopf/acopf_generator_kernel_ka.jl
@@ -29,12 +29,11 @@ function generator_kernel_two_level(
     zu, lu, rho_u, device
 )
     nblk = div(model.grid_data.ngen, 32, RoundUp)
-    ev = generator_kernel_two_level_ka(device,32,nblk*32)(
+    generator_kernel_two_level_ka(device,32,nblk*32)(
         baseMVA, model.grid_data.ngen, model.gen_start,
         u, xbar, zu, lu, rho_u, model.pgmin_curr, model.pgmax_curr, model.grid_data.qgmin, model.grid_data.qgmax,
-        model.grid_data.c2, model.grid_data.c1,
-        dependencies=Event(device)
+        model.grid_data.c2, model.grid_data.c1
     )
-    wait(ev)
+    KA.synchronize(device)
     return 0.0
 end
diff --git a/src/models/acopf/acopf_init_solution_ka.jl b/src/models/acopf/acopf_init_solution_ka.jl
@@ -60,18 +60,16 @@ function init_solution!(
 
     data = model.grid_data
 
-    ev = init_generator_kernel_one_level_ka(device, 64, data.ngen)(
+    init_generator_kernel_one_level_ka(device, 64, data.ngen)(
         data.ngen, model.gen_start,
-        data.pgmax, data.pgmin, data.qgmax, data.qgmin, sol.v_curr,
-        dependencies=Event(device)
+        data.pgmax, data.pgmin, data.qgmax, data.qgmin, sol.v_curr
     )
-    wait(ev)
-    ev = init_branch_bus_kernel_one_level_ka(device, 64, data.nline)(
+    KA.synchronize(device)
+    init_branch_bus_kernel_one_level_ka(device, 64, data.nline)(
         data.nline, model.line_start, rho_va,
         data.brBusIdx, data.Vmax, data.Vmin, data.YffR, data.YffI, data.YftR, data.YftI,
-        data.YtfR, data.YtfI, data.YttR, data.YttI, sol.u_curr, sol.v_curr, sol.rho,
-        dependencies=Event(device)
+        data.YtfR, data.YtfI, data.YttR, data.YttI, sol.u_curr, sol.v_curr, sol.rho
     )
-    wait(ev)
+    KA.synchronize(device)
 
 end
diff --git a/src/models/acopf/acopf_set_linelimit_ka.jl b/src/models/acopf/acopf_set_linelimit_ka.jl
@@ -17,5 +17,5 @@ function acopf_set_linelimit(
     device
 )
     ev = set_rateA_kernel_ka(device,64,mod.nline)(mod.nline, mod.membuf, mod.rateA)
-    wait(ev)
+    KA.synchronize(device)
 end
diff --git a/src/utils/utilities_ka.jl b/src/utils/utilities_ka.jl
@@ -39,7 +39,8 @@ end
 function LinearAlgebra.norm(x, device)
     y = KAArray{Float64}(1, device)
     n = length(x)
-    wait(norm_kernel(device)(Val{n}(), x, y, ndrange=n, dependencies=Event(device)))
+    norm_kernel(device)(Val{n}(), x, y, ndrange=n)
+    KA.synchronize(device)
     ret = y |> Array
     return ret[1]
 end
diff --git a/test/Project.toml b/test/Project.toml
@@ -1,11 +1,9 @@
 [deps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-CUDAKernels = "72cfdca4-0801-4ab0-bf6a-d52aa10adc57"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
-ROCKernels = "7eb9e9f0-4bd3-4c4c-8bef-26bd9629d9b9"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/algorithms/acopf_update_ka.jl b/test/algorithms/acopf_update_ka.jl
@@ -5,20 +5,18 @@ KA = KernelAbstractions
 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}
+       function ExaAdmm.KAArray{T}(n::Int, device::CUDABackend) where {T}
            return CuArray{T}(undef, n)
        end
-       push!(devices, CUDADevice())
+       push!(devices, CUDABackend())
    end
     if AMDGPU.has_rocm_gpu()
-        using ROCKernels
         # Set for crusher login node to avoid other users
         AMDGPU.default_device!(AMDGPU.devices()[2])
-        function ExaAdmm.KAArray{T}(n::Int, device::ROCDevice) where {T}
+        function ExaAdmm.KAArray{T}(n::Int, device::ROCBackend) where {T}
             return ROCArray{T}(undef, n)
         end
-        push!(devices, ROCDevice())
+        push!(devices, ROCBackend())
     end
 end
 @testset "Testing [x,xbar,z,l,lz] updates and a solve for ACOPF using KA" for _device in devices