Initial wrap of hipTENSOR library

docs/src/index.md

@@ -105,6 +105,7 @@ gridsize = cld(length(c), groupsize)
 | FFT | `fft`/`plan_fft` via rocFFT | [Fourier Transforms](libraries/fft.md) |
 | Random | `rand`/`randn` via rocRAND | [Random Numbers](libraries/rand.md) |
 | Deep learning | Conv/pooling/softmax via MIOpen (NNlib/Flux) | [Deep Learning (MIOpen)](libraries/dnn.md) |
+| Tensor operations | contract/permute/reduce via hipTensor | [Tensor Operations](libraries/tensor.md) |
 
 ## Questions and Contributions
 

docs/src/libraries/tensor.md

@@ -0,0 +1,10 @@
+# Tensor Operations
+
+`AMDGPU.jl` now provides wrappers for the [`hipTensor`](https://rocm.docs.amd.com/projects/hipTensor/en/latest/index.html) library,
+which can be used to perform operations on high-dimensional arrays. However, there are some caveats:
+
+- `hipTensor` isn't supported on every AMD GPU. You can find the list of supported GPUs [at the `hipTensor` documentation](https://rocm.docs.amd.com/projects/hipTensor/en/latest/api-reference/api-reference.html#supported-gpu-architectures).
+- for [`hiptensorCreateElementwiseTrinary`](https://rocm.docs.amd.com/projects/hipTensor/en/latest/api-reference/api-reference.html#hiptensorcreateelementwisetrinary), `hipTENSOR` 2.2 has the two binary operators the wrong way round: it applies the
+  operator passed as `opABC` to `A` and `B`, and the one passed as `opAB` to that intermediate result and `C`. This is currently handled by our wrapper.
+- Unlike NVIDIA's `cuTENSOR`, `hipTENSOR` wants a *complex* compute descriptor for complex operands: pairing e.g. `ComplexF32` tensors with `HIPTENSOR_COMPUTE_DESC_32F` makes
+  `hiptensorCreatePlan` fail with `HIPTENSOR_STATUS_EXECUTION_FAILED`

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

@@ -128,6 +128,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 AMDGPU_LLVM_Backend_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,59 @@
+module hipTENSOR
+
+using AMDGPU
+using GPUToolbox: @gcsafe_ccall, @checked, @enum_without_prefix, @debug_ccall
+using AMDGPU: Mem
+
+using BFloat16s: BFloat16
+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("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 = hiptensorGetVersion()
+    major = ver ÷ 10000
+    minor = (ver ÷ 100) % 100
+    patch = ver % 100
+    return VersionNumber(join((major,minor,patch), "."))
+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