hip/rocm - build and use pufferlib cuda extension

[sozforex]

Enables compilation and use of cuda extension on hip/rocm systems.

Requires installing through uv pip install --no-build-isolation or running python setup.py build_ext --inplace.

There should be a better way to detect if the cuda extension exists and should be used in pufferl.py.

I've only tested this with an amd gpu.

[elevatorguy]

Is it really that simple for AMD? (INTEL_HOME would rely upon XPU pytorch support - likely /opt/intel/oneapi in Linux if using Intel's compiler)

[sozforex]

For amd this does rely on torch being compiled with rocm, and torch being able to find rocm in the system (either through ROCM_HOME or ROCM_PATH env variable, or through heuristics based on existence/location of hipcc).

torch.utils.cpp_extension.ROCM_HOME [or torch.utils.cpp_extension.CUDA_HOME] is populated only if torch is compiled with rocm [or cuda for CUDA_HOME] support,
AND if torch is able to find corresponding env variable set or find the the path to ROCM_HOME through some heuristics [through existence/location of hipcc]

For amd:
https://github.com/pytorch/pytorch/blob/5743d731c1de495ecf3bb03682a2dcbe207ca895/torch/utils/cpp_extension.py#L229
https://github.com/pytorch/pytorch/blob/5743d731c1de495ecf3bb03682a2dcbe207ca895/torch/utils/cpp_extension.py#L121

For cuda:
https://github.com/pytorch/pytorch/blob/5743d731c1de495ecf3bb03682a2dcbe207ca895/torch/utils/cpp_extension.py#L236
https://github.com/pytorch/pytorch/blob/5743d731c1de495ecf3bb03682a2dcbe207ca895/torch/utils/cpp_extension.py#L95

I'm not really familiar with intel ecosystem, I guess that there is no ~ cuda-compatible intel equivalent to hip/rocm.

pytorch automatically translates cuda extensions using hipify script such that they can be compiled with hipcc and used with amd gpus, though there can be various problems, and in some cases it may be required to manually edit/fix the generated code.

[elevatorguy]

icpx is Intel's C++ (icx for C) compile command; thanks for pointing out intel-sycl-rt - hadn't seen torch/utils/cpp_extension.py yet.

[elevatorguy]

Not sure why SYCL_HOME instead of INTEL_HOME though the former may have been chosen for reason(s); subjective or otherwise.

Is it really that simple for AMD? (INTEL_HOME would rely upon XPU pytorch support - likely /opt/intel/oneapi in Linux if using Intel's compiler)

Will try using SYCL_HOME.