Initial wrap of hipTENSOR library

gen/hiptensor/generator.jl

@@ -0,0 +1,21 @@
+using Clang.Generators
+using JuliaFormatter
+
+include_dir = normpath("/opt/rocm/include")
+rocblas_dir = joinpath(include_dir, "hiptensor")
+options = load_options("hiptensor/hiptensor-generator.toml")
+
+args = get_default_args()
+push!(args, "-I$include_dir")
+
+headers = [
+    joinpath(rocblas_dir, header)
+    for header in readdir(rocblas_dir)
+    if endswith(header, ".h")
+]
+
+ctx = create_context(headers, args, options)
+build!(ctx)
+
+path = options["general"]["output_file_path"]
+format_file(path, YASStyle())

gen/hiptensor/hiptensor-generator.toml

@@ -0,0 +1,8 @@
+[general]
+library_name = "libhiptensor"
+output_file_path = "./libhiptensor.jl"
+export_symbol_prefixes = []
+print_using_CEnum = false
+
+[codegen]
+use_ccall_macro = true

src/AMDGPU.jl

@@ -4,6 +4,8 @@ using Adapt
 using BFloat16s
 using CEnum
 using GPUCompiler
+using GPUToolbox
+using GPUToolbox: @enum_without_prefix
 using GPUArrays
 using GPUArrays: allowscalar
 using Libdl
@@ -114,6 +116,7 @@ include("kernels/reverse.jl")
 include("blas/rocBLAS.jl")
 include("solver/rocSOLVER.jl")
 include("sparse/rocSPARSE.jl")
+include("tensor/hipTENSOR.jl")
 include("rand/rocRAND.jl")
 include("fft/rocFFT.jl")
 include("dnn/MIOpen.jl")

src/discovery/discovery.jl

@@ -3,6 +3,7 @@ module ROCmDiscovery
 export lld_artifact, lld_path, libhsaruntime, libdevice_libs, libhip
 export librocblas, librocsparse, librocsolver
 export librocrand, librocfft, libMIOpen_path
+export libhiptensor
 
 using LLD_jll
 using ROCmDeviceLibs_jll
@@ -55,6 +56,7 @@ global librocsparse::String = ""
 global librocsolver::String = ""
 global librocrand::String = ""
 global librocfft::String = ""
+global libhiptensor::String = ""
 global libMIOpen_path::String = ""
 
 function __init__()
@@ -96,6 +98,7 @@ function __init__()
         global librocsolver = find_rocm_library(lib_prefix * "rocsolver"; rocm_path)
         global librocrand = find_rocm_library(lib_prefix * "rocrand"; rocm_path)
         global librocfft = find_rocm_library(lib_prefix * "rocfft"; rocm_path)
+        global libhiptensor = find_rocm_library(lib_prefix * "hiptensor"; rocm_path)
         global libMIOpen_path = find_rocm_library(lib_prefix * "MIOpen"; rocm_path)
     catch err
         @error """ROCm discovery failed!

src/tensor/error.jl

@@ -0,0 +1,46 @@
+export hipTENSORError
+
+struct hipTENSORError <: Exception
+    code::hiptensorStatus_t
+end
+
+Base.convert(::Type{hiptensorStatus_t}, err::hipTENSORError) = err.code
+
+Base.showerror(io::IO, err::hipTENSORError) =
+    print(io, "hipTENSORError: ", description(err), " (code $(reinterpret(Int32, err.code)), $(name(err)))")
+
+name(err::hipTENSORError) = unsafe_string(hiptensorGetErrorString(err))
+
+## COV_EXCL_START
+function description(err::hipTENSORError)
+    if err.code == HIPTENSOR_STATUS_SUCCESS
+        "the operation completed successfully"
+    elseif err.code == HIPTENSOR_STATUS_NOT_INITIALIZED
+        "the library was not initialized"
+    elseif err.code == HIPTENSOR_STATUS_ALLOC_FAILED
+        "the resource allocation failed"
+    elseif err.code == HIPTENSOR_STATUS_INVALID_VALUE
+        "an invalid value was used as an argument"
+    elseif err.code == HIPTENSOR_STATUS_ARCH_MISMATCH
+        "an absent device architectural feature is required"
+    elseif err.code == HIPTENSOR_STATUS_EXECUTION_FAILED
+        "the GPU program failed to execute"
+    elseif err.code == HIPTENSOR_STATUS_INTERNAL_ERROR
+        "an internal operation failed"
+    elseif err.code == HIPTENSOR_STATUS_NOT_SUPPORTED
+        "operation not supported (yet)"
+    elseif err.code == HIPTENSOR_STATUS_CK_ERROR
+        "error detected trying to check the license"
+    elseif err.code == HIPTENSOR_STATUS_HIP_ERROR
+        "error occurred during a HIP operation"
+    elseif err.code == HIPTENSOR_STATUS_INSUFFICIENT_WORKSPACE
+        "insufficient workspace memory for this operation"
+    elseif err.code == HIPTENSOR_STATUS_INSUFFICIENT_DRIVER
+        "insufficient driver version"
+    elseif err.code == HIPTENSOR_STATUS_IO_ERROR
+        "file not found"
+    else
+        "no description for this error"
+    end
+end
+## COV_EXCL_STOP

src/tensor/hipTENSOR.jl

@@ -0,0 +1,60 @@
+module hipTENSOR
+
+using AMDGPU
+using AMDGPU
+using AMDGPU: @gcsafe_ccall, @checked, @enum_without_prefix, @debug_ccall
+using AMDGPU: Mem
+import AMDGPU: libhiptensor, HandleCache, HIP, library_state
+import AMDGPU.Mem: alloc_or_retry!
+import .HIP: HIPContext, HIPStream, hipStream_t
+
+using CEnum: @cenum
+
+using Printf: @printf
+
+export has_hiptensor
+
+has_hiptensor() = AMDGPU.functional(:hiptensor)
+
+# core library
+include("libhiptensor.jl")
+
+# low-level wrappers
+include("error.jl")
+#include("utils.jl")
+include("types.jl")
+include("operations.jl")
+
+# high-level integrations
+include("interfaces.jl")
+
+
+## handles
+
+function create_handle()
+    AMDGPU.functional(:hiptensor) || error("hipTENSOR is not available")
+
+    handle_ref = Ref{hiptensorHandle_t}()
+    hiptensorCreate(handle_ref)
+    handle_ref[]
+end
+
+const IDLE_HANDLES = HandleCache{HIPContext,hiptensorHandle_t}()
+
+lib_state() = library_state(
+    :hipTENSOR, hiptensorHandle_t, IDLE_HANDLES,
+    create_handle, hiptensorDestroy, (handle, stream) -> 0 )
+
+handle() = lib_state().handle
+stream() = lib_state().stream
+
+function version()
+    #=ver = Ref{Cint}()
+    hiptensor_get_version(handle(), ver)
+    major = ver[] ÷ 100000
+    minor = (ver[] ÷ 100) % 1000
+    patch = ver[] % 100=#
+    return VersionNumber(2, 2, 0)
+end
+
+end

src/tensor/interfaces.jl

@@ -0,0 +1,62 @@
+# interfacing with other packages
+
+## Base
+
+function Base.:(+)(A::hipTensor, B::hipTensor)
+    α = convert(eltype(A), 1.0)
+    γ = convert(eltype(B), 1.0)
+    C = similar(B)
+    elementwise_binary_execute!(α, A.data, A.inds, HIPTENSOR_OP_IDENTITY,
+                                γ, B.data, B.inds, HIPTENSOR_OP_IDENTITY,
+                                C.data, C.inds, HIPTENSOR_OP_ADD)
+    return C
+end
+
+function Base.:(-)(A::hipTensor, B::hipTensor)
+    α = convert(eltype(A), 1.0)
+    γ = convert(eltype(B), -1.0)
+    C = similar(B)
+    elementwise_binary_execute!(α, A.data, A.inds, HIPTENSOR_OP_IDENTITY,
+                                γ, B.data, B.inds, HIPTENSOR_OP_IDENTITY,
+                                C.data, C.inds, HIPTENSOR_OP_ADD)
+    return C
+end
+
+function Base.:(*)(A::hipTensor, B::hipTensor)
+    tC = promote_type(eltype(A), eltype(B))
+    A_uniqs = [(idx, i) for (idx, i) in enumerate(A.inds) if !(i in B.inds)]
+    B_uniqs = [(idx, i) for (idx, i) in enumerate(B.inds) if !(i in A.inds)]
+    A_sizes = map(x->size(A,x[1]), A_uniqs)
+    B_sizes = map(x->size(B,x[1]), B_uniqs)
+    A_inds = map(x->x[2], A_uniqs)
+    B_inds = map(x->x[2], B_uniqs)
+    C = hipTensor(fill!(similar(B.data, tC, Dims(vcat(A_sizes, B_sizes))), zero(tC)), vcat(A_inds, B_inds))
+    return mul!(C, A, B)
+end
+
+
+## LinearAlgebra
+
+using LinearAlgebra
+
+function LinearAlgebra.axpy!(a, X::hipTensor, Y::hipTensor)
+    elementwise_binary_execute!(a, X.data, X.inds, HIPTENSOR_OP_IDENTITY,
+                                one(eltype(Y)), Y.data, Y.inds, HIPTENSOR_OP_IDENTITY,
+                                Y.data, Y.inds, HIPTENSOR_OP_ADD)
+    return Y
+end
+
+function LinearAlgebra.axpby!(a, X::hipTensor, b, Y::hipTensor)
+    elementwise_binary_execute!(a, X.data, X.inds, HIPTENSOR_OP_IDENTITY,
+                                b, Y.data, Y.inds, HIPTENSOR_OP_IDENTITY,
+                                Y.data, Y.inds, HIPTENSOR_OP_ADD)
+    return Y
+end
+
+function LinearAlgebra.mul!(C::hipTensor, A::hipTensor, B::hipTensor, α::Number, β::Number)
+   contract!(α, A.data, A.inds, HIPTENSOR_OP_IDENTITY,
+             B.data, B.inds, HIPTENSOR_OP_IDENTITY, β,
+             C.data, C.inds, HIPTENSOR_OP_IDENTITY, HIPTENSOR_OP_IDENTITY;
+             jit=HIPTENSOR_JIT_MODE_DEFAULT)
+   return C
+end