Fixing norm for AMD

Project.toml

@@ -18,7 +18,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-AMDGPU = "=0.4.2"
+AMDGPU = "0.4"
 CUDA = "3.4"
 ExaTron = "2"
 FileIO = "1.14"

src/interface/solve_acopf.jl

@@ -17,14 +17,12 @@ function solve_acopf(case::String;
     elseif use_gpu && isa(ka_device, Nothing)
         CUDA.device!(gpu_no)
         TD = CuArray{Float64,1}; TI = CuArray{Int,1}; TM = CuArray{Float64,2}
-    elseif use_gpu && isa(ka_device, KA.Device)
-        if has_cuda_gpu()
+    elseif has_cuda_gpu()
             TD = CuArray{Float64,1}; TI = CuArray{Int,1}; TM = CuArray{Float64,2}
-        elseif has_rocm_gpu()
-            TD = ROCArray{Float64,1}; TI = ROCArray{Int,1}; TM = ROCArray{Float64,2}
-        end
+    elseif has_rocm_gpu()
+        TD = ROCArray{Float64,1}; TI = ROCArray{Int,1}; TM = ROCArray{Float64,2}
     else
-        error("Inconsistent device selection use_gpu=$use_gpu and ka_device=$(typepof(ka_device))")
+        error("Inconsistent device selection use_gpu=$use_gpu and ka_device=$(typeof(ka_device))")
     end
 
     env = AdmmEnv{T,TD,TI,TM}(case, rho_pq, rho_va; case_format=case_format,

src/models/acopf/acopf_admm_prepoststep_ka.jl

@@ -4,7 +4,7 @@ function admm_outer_prestep(
     device
 )
     sol, info = mod.solution, mod.info
-    info.norm_z_prev = norm(sol.z_curr)
+    info.norm_z_prev = norm(sol.z_curr, device)
     return
 end
 

src/models/acopf/acopf_admm_update_residual_ka.jl

@@ -33,10 +33,10 @@ function admm_update_residual(
     )
     wait(ev)
 
-    info.primres = norm(sol.rp)
-    info.dualres = norm(sol.rd)
-    info.norm_z_curr = norm(sol.z_curr)
-    info.mismatch = norm(sol.Ax_plus_By)
+    info.primres = norm(sol.rp, device)
+    info.dualres = norm(sol.rd, device)
+    info.norm_z_curr = norm(sol.z_curr, device)
+    info.mismatch = norm(sol.Ax_plus_By, device)
 
     return
 end

src/utils/utilities_ka.jl

@@ -17,3 +17,27 @@ end
         c[tx] = a[tx] - b[tx]
     end
 end
+@kernel function norm_kernel(::Val{n}, x,
+                    y
+                    ) where {n}
+    tx = @index(Local, Linear)
+    bx = @index(Group, Linear)
+
+    _x = @localmem Float64 (n,)
+    _x[tx] = x[tx]
+    @synchronize
+
+    v = ExaTron.dnrm2(n, _x, 1, tx)
+    if bx == 1 && tx == 1
+        y[1] = v
+    end
+    @synchronize
+
+end
+function LinearAlgebra.norm(x, device)
+    y = KAArray{Float64}(1, device)
+    n = length(x)
+    wait(norm_kernel(device,n)(Val{n}(), x, y,ndrange=(n,),dependencies=Event(device)))
+    ret = y |> Array
+    return ret[1]
+end

test/algorithms/acopf_update_cpu.jl

@@ -176,5 +176,5 @@
     @test mod.info.status == :Solved
     @test mod.info.outer == 20
     @test mod.info.cumul == 1232
-    @test isapprox(mod.info.objval, 129645.676; atol=1e-3)
+    @test isapprox(mod.info.objval, 129645.676; rtol=1e-6)
 end

test/algorithms/acopf_update_ka.jl

@@ -13,6 +13,8 @@ if CUDA.has_cuda_gpu() || AMDGPU.has_rocm_gpu()
    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}
             return ROCArray{T}(undef, n)
         end
@@ -31,7 +33,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, ROCDevice)
+    else
         if CUDA.has_cuda_gpu()
             TD = CuArray{Float64,1}; TI = CuArray{Int,1}; TM = CuArray{Float64,2}
         elseif AMDGPU.has_rocm_gpu()