Cleanup & show shmem size for devices

Author: pxl-th

src/AMDGPU.jl

@@ -154,9 +154,8 @@ import .ROCKernels: ROCBackend
 export ROCBackend
 
 function __init__()
-    atexit() do
-        Runtime.RT_EXITING[] = true
-    end
+    # Used to shutdown hostcalls if any is running.
+    atexit(() -> begin Runtime.RT_EXITING[] = true end)
 
     if haskey(ENV, "HIP_LAUNCH_BLOCKING")
         launch_blocking = parse(Bool, ENV["HIP_LAUNCH_BLOCKING"])
@@ -167,78 +166,57 @@ function __init__()
         end
     end
 
-    # Quiet path first, in case this system doesn't have AMD GPUs
-    if Sys.islinux() && !ispath("/dev/kfd")
-        @debug "/dev/kfd not available (no AMD GPU), skipping initialization"
-        return
-    end
-
-    # Verbose path, something is misconfigured
     if Sys.islinux()
+        if !ispath("/dev/kfd")
+            @debug "/dev/kfd not available (no AMD GPU), skipping initialization"
+            return
+        end
+
         if !isempty(libhsaruntime)
-            # Initialize the HSA runtime.
-            status = HSA.init()
-            if status == HSA.STATUS_SUCCESS
-                # Register shutdown hook.
-                atexit(() -> HSA.shut_down())
-            else
+            HSA.init() == HSA.STATUS_SUCCESS ?
+                atexit(() -> HSA.shut_down()) :
                 @warn "HSA initialization failed with code $status"
-            end
         else
-            @warn """
-            HSA runtime is unavailable, compilation and runtime functionality will be disabled.
-            """
+            @warn "HSA runtime is unavailable, compilation and runtime functionality will be disabled."
             if parse(Bool, get(ENV, "JULIA_AMDGPU_CORE_MUST_LOAD", "0"))
                 print_build_diagnostics()
                 error("Failed to load HSA runtime, but HSA must load, bailing out")
             end
         end
     end
 
-    # Check whether ld.lld was found
     if !functional(:lld)
-        @warn """
-        LLD is unavailable, compilation functionality will be disabled.
-        """
+        @warn "LLD is unavailable, compilation functionality will be disabled."
         if parse(Bool, get(ENV, "JULIA_AMDGPU_CORE_MUST_LOAD", "0"))
             print_build_diagnostics()
             error("Failed to find ld.lld, but ld.lld must exist, bailing out")
         end
     end
 
-    # Check whether device intrinsics are available
     if !functional(:device_libs)
-        @warn """
-        Device libraries are unavailable, device intrinsics will be disabled.
-        """
+        @warn "Device libraries are unavailable, device intrinsics will be disabled."
         if parse(Bool, get(ENV, "JULIA_AMDGPU_CORE_MUST_LOAD", "0"))
             print_build_diagnostics()
             error("Failed to find Device Libs, but Device Libs must exist, bailing out")
         end
     end
 
-    # Check whether HIP is available
     if functional(:hip)
         HIP.devices()
     else
-        @warn """
-        HIP library is unavailable, HIP integration will be disabled.
-        """
+        @warn "HIP library is unavailable, HIP integration will be disabled."
         if parse(Bool, get(ENV, "JULIA_AMDGPU_HIP_MUST_LOAD", "0"))
             print_build_diagnostics()
             error("Failed to load HIP runtime, but HIP must load, bailing out")
         end
     end
 
-    # Check whether external libraries are available.
     hiplibs = (
         ("rocBLAS", :rocblas), ("rocSPARSE", :rocsparse),
         ("rocSOLVER", :rocsolver), ("rocALUTION", :rocalution),
         ("rocRAND", :rocrand), ("rocFFT", :rocfft), ("MIOpen", :MIOpen))
     for (name, symbol) in hiplibs
-        if !functional(symbol)
-            @warn "$name is unavailable, functionality will be disabled."
-        end
+        functional(symbol) || @warn "$name is unavailable, functionality will be disabled."
     end
 end
 

src/hip/device.jl

@@ -161,17 +161,18 @@ function __pretty_data(dev::HIPDevice)
     name_ptr = pointer([props.name...])
     name = unsafe_string(name_ptr)
     reshape(String[
-        "$(dev.device_id)", name, "$(dev.gcn_arch)",
-        "$(dev.wavefrontsize)", "$(Base.format_bytes(props.totalGlobalMem))",
+        "$(dev.device_id)", name, dev.gcn_arch, "$(dev.wavefrontsize)",
+        Base.format_bytes(props.totalGlobalMem),
+        Base.format_bytes(props.sharedMemPerBlock),
     ], 1, :)
 end
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, dev::HIPDevice)
     PrettyTables.pretty_table(io, __pretty_data(dev); header=[
-        "Id", "Name", "GCN arch", "Wavefront", "Memory"])
+        "Id", "Name", "GCN arch", "Wavefront", "Memory", "Shared Memory"])
 end
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, devs::Vector{HIPDevice})
     PrettyTables.pretty_table(io, vcat(__pretty_data.(devs)...); header=[
-        "Id", "Name", "GCN arch", "Wavefront", "Memory"])
+        "Id", "Name", "GCN arch", "Wavefront", "Memory", "Shared Memory"])
 end