Initial wrap of hipTENSOR library

[kshyatt]

Mostly copied from cuTENSOR. Tests are failing with weird segfaults I don't understand, I'll try to keep debugging, but if anyone else wants to drop in and help I'd be very appreciative as I am quite unfamiliar with HIP.

[luraess]

cscs-ci run

[luraess]

Thanks for pushing this through, the wrapper reads well and the notes left about hipTENSOR's real behaviour are useful.

• The arch gate seems the blocker. hipTENSOR is CDNA only (Instinct accelerators) and exits the process on load on RDNA (Radeon, the Buildkite CI GPUs), while functional(:hiptensor) is a presence-only check, so on our runners this kills the test run instead of skipping it. Something we have to gate around it. The version() stub is the other one we should fix before merging.

• I'd reword some of the comments. The 2.2 quirk notes are load-bearing: the mode-order limitation is why reordered_descriptor exists, the swapped opAB/opABC explains a deliberate argument swap, and the "the plan does not copy what it is built from" note is why hipTensorPlan carries refs. Four of them are phrased against cuTENSOR, which is not needed and they could stand on their own.

Maybe the cuTENSOR deltas and the version-specific quirks should go in some succinct doc. There are no docs in this PR and docs/src/libraries/ already has dnn/fft/linalg/rand/sparse, so a tensor.md with a "Known hipTENSOR 2.2 limitations" section would fit nicely, with short pointers at the call sites (also, easier to re-verify when ROCm bumps).

Also, :hiptensor is added to the functional docstring and dispatch but not to the :all tuple. Right or oversight?

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: 241c7de	Previous: 460b5d4	Ratio
amdgpu/synchronization/context/device	550 ns	552.75 ns	1.00
amdgpu/synchronization/stream/blocking	230 ns	232.5 ns	0.99
amdgpu/synchronization/stream/nonblocking	312.5 ns	315 ns	0.99
array/accumulate/Float32/1d	77341.25 ns	80761 ns	0.96
array/accumulate/Float32/dims=1	270336.5 ns	263598.75 ns	1.03
array/accumulate/Float32/dims=1L	81656.25 ns	91826.5 ns	0.89
array/accumulate/Float32/dims=2	80681.25 ns	87378.75 ns	0.92
array/accumulate/Float32/dims=2L	2757224.25 ns	2753302.75 ns	1.00
array/accumulate/Int64/1d	80606 ns	81731 ns	0.99
array/accumulate/Int64/dims=1	245393.25 ns	247055.75 ns	0.99
array/accumulate/Int64/dims=1L	84101.25 ns	97598.75 ns	0.86
array/accumulate/Int64/dims=2	86408.5 ns	90248.75 ns	0.96
array/accumulate/Int64/dims=2L	2897791.25 ns	2891959.5 ns	1.00
array/broadcast	73991 ns	68128.25 ns	1.09
array/construct	2092.5 ns	2475 ns	0.85
array/copy	37403 ns	37450.5 ns	1.00
array/copyto!/cpu_to_gpu	111431.5 ns	111219 ns	1.00
array/copyto!/gpu_to_cpu	111809 ns	112099 ns	1.00
array/copyto!/gpu_to_gpu	48668.25 ns	47263.25 ns	1.03
array/iteration/findall/bool	138882 ns	139007 ns	1.00
array/iteration/findall/int	151929.75 ns	151309.5 ns	1.00
array/iteration/findfirst/bool	147452 ns	145502 ns	1.01
array/iteration/findfirst/int	146832 ns	147064.5 ns	1.00
array/iteration/findmin/1d	121699.25 ns	128964.25 ns	0.94
array/iteration/findmin/2d	112834 ns	116704 ns	0.97
array/iteration/logical	249168.5 ns	253920.75 ns	0.98
array/iteration/scalar	309006.75 ns	306916.5 ns	1.01
array/permutedims/2d	72586.25 ns	72898.5 ns	1.00
array/permutedims/3d	71743.5 ns	72343.5 ns	0.99
array/permutedims/4d	74721.25 ns	75161.25 ns	0.99
array/random/rand/Float32	45880.75 ns	45988 ns	1.00
array/random/rand/Int64	55018.25 ns	55215.75 ns	1.00
array/random/rand!/Float32	66158.5 ns	66023.25 ns	1.00
array/random/rand!/Int64	73886 ns	67061 ns	1.10
array/random/randn/Float32	79521 ns	81601 ns	0.97
array/random/randn!/Float32	72591.25 ns	81943.75 ns	0.89
array/reductions/mapreduce/Float32/1d	101116.5 ns	105458.75 ns	0.96
array/reductions/mapreduce/Float32/dims=1	86363.75 ns	95684 ns	0.90
array/reductions/mapreduce/Float32/dims=1L	832366.75 ns	833506.5 ns	1.00
array/reductions/mapreduce/Float32/dims=2	88871.25 ns	96211.5 ns	0.92
array/reductions/mapreduce/Float32/dims=2L	144147 ns	144927 ns	0.99
array/reductions/mapreduce/Int64/1d	101764 ns	107101.5 ns	0.95
array/reductions/mapreduce/Int64/dims=1	91743.75 ns	92111.25 ns	1.00
array/reductions/mapreduce/Int64/dims=1L	831696.75 ns	832748.75 ns	1.00
array/reductions/mapreduce/Int64/dims=2	88423.75 ns	95113.75 ns	0.93
array/reductions/mapreduce/Int64/dims=2L	145119.5 ns	145554.5 ns	1.00
array/reductions/reduce/Float32/1d	101053.75 ns	105439 ns	0.96
array/reductions/reduce/Float32/dims=1	88551.25 ns	95611.25 ns	0.93
array/reductions/reduce/Float32/dims=1L	837009.5 ns	832846.5 ns	1.00
array/reductions/reduce/Float32/dims=2	87506.25 ns	96106.5 ns	0.91
array/reductions/reduce/Float32/dims=2L	144207.25 ns	144737 ns	1.00
array/reductions/reduce/Int64/1d	102386.5 ns	106451.25 ns	0.96
array/reductions/reduce/Int64/dims=1	91033.75 ns	95993.75 ns	0.95
array/reductions/reduce/Int64/dims=1L	833424.25 ns	833469 ns	1.00
array/reductions/reduce/Int64/dims=2	88279 ns	95636.5 ns	0.92
array/reductions/reduce/Int64/dims=2L	145187.25 ns	145059.5 ns	1.00
array/reverse/1d	46353.25 ns	46908.25 ns	0.99
array/reverse/1dL	76573.5 ns	76671 ns	1.00
array/reverse/1dL_inplace	80996.25 ns	81231.25 ns	1.00
array/reverse/1d_inplace	61908.5 ns	53403 ns	1.16
array/reverse/2d	50975.75 ns	50923.25 ns	1.00
array/reverse/2dL	86163.75 ns	88333.75 ns	0.98
array/reverse/2dL_inplace	92993.75 ns	92821.25 ns	1.00
array/reverse/2d_inplace	46398.25 ns	56983.25 ns	0.81
array/sorting/1d	333667.25 ns	335852 ns	0.99
integration/byval/reference	40081 ns	39781 ns	1.01
integration/byval/slices=1	40841 ns	41850 ns	0.98
integration/byval/slices=2	148192 ns	160152 ns	0.93
integration/byval/slices=3	237914 ns	238584 ns	1.00
integration/volumerhs	4986822 ns	4888917 ns	1.02
kernel/indexing	58288.25 ns	58573.5 ns	1.00
kernel/indexing_checked	60360.75 ns	60428.5 ns	1.00
kernel/launch	1407.5 ns	1425 ns	0.99
kernel/rand	103184 ns	84416 ns	1.22
latency/import	1681739383 ns	1683857388 ns	1.00
latency/precompile	38490964017 ns	37917346182 ns	1.02
latency/ttfp	5733669451 ns	5744789967 ns	1.00

This comment was automatically generated by workflow using github-action-benchmark.

[kshyatt]

@luraess did the updates satisfy the concern you had?

[luraess]

ROCm ships this exact list as hiptensorSupportedArchitectures.cmake in the hipTENSOR cmake dir, so parsing it when present (falling back to this tuple for runtime-only installs) would track the install rather than drift each release. Could be worth doing because gfx940/gfx941 make the constant a superset of what 7.x actually supports, and superset may be unsafe as those devices would pass the gate and we would be back to dlopen deciding the outcome.

[luraess]

Thanks. I left a few minor things but looks good. The one remaining is whether you'd be willing to add something to the docs, like docs/src/libraries/tensor.md. The in file comments on porting notes and 2.2 quirks are still spread across three files and having these listed into the docs could be helpful, maybe.

[kshyatt]

Added some doc notes :)

[luraess]

why is CI crying red... Seems missing a HIP import on MI300 and idk what BK is up to...

[kshyatt]

Let me rebase on top of main and see if it was intermittent?

[luraess]

LGTM now.

[luraess]

Thanks for the addition!