Add support for AMD GPUs

As some exploratory work, we should build an image for ROCm and see if it works.