Add AMD GPU support via the HIP backend

.gitignore

@@ -1,5 +1,13 @@
 *.swp
 *.swo
+# PTX (CUDA backend) and HIP code objects (HIP backend) are build artifacts
+# whose bytes are embedded into the generated *_wrapper.go / *_wrapper_hip.go
+# files; do not track them, except the small fixtures the cu module test loads.
+*.ptx
+!cuda/cu/testdata/*.ptx
+cuda/*.co
+!cuda/cu/testdata/*.co
+cuda/cuda2go
 *.5
 *.6
 *.8

README.md

@@ -24,6 +24,8 @@ Consider downloading a [pre-compiled mumax³ binary](https://mumax.github.io/dow
 
 If you want to compile nevertheless, 4 essential components will be required to build mumax³: an ***NVIDIA driver***, ***Go***, ***CUDA*** (≤12.9) and ***C***.
 
+mumax³ also runs on ***AMD GPUs*** through an opt-in HIP backend: install ***ROCm*** in place of the NVIDIA driver and CUDA toolkit (***Go*** and ***C*** are still needed), then see *Build for AMD GPUs* below.
+
 * *If they are not yet present on your system*: install them as detailed below.
 * *If they are already installed*: check if they work correctly by running the *check* for each component written below.
 
@@ -118,6 +120,26 @@ Click on the arrows below to expand the installation instructions:<br><sub><sup>
 
 </details>
 
+<details><summary><b><i>Build for AMD GPUs</i></b> (ROCm/HIP, instead of CUDA)</summary>
+
+mumax³ also runs on AMD GPUs through an additive, opt-in HIP backend; the CUDA path is unchanged and remains the default. Instead of the NVIDIA driver and CUDA toolkit, install [ROCm](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/) (the HIP runtime and headers; ROCm 7.2 or newer is recommended). ***Go*** and a ***C*** compiler are still required, as above.
+
+Once ROCm, Go and C are in place, follow the *Building from source* steps below to clone the repository and initialise the Go module, then build the HIP backend instead of running `make`:
+
+```bash
+go install -tags hip github.com/mumax/3/...
+```
+
+The generated HIP wrappers, with the device code already embedded, are committed, so this command needs only the ROCm runtime (no `hipcc`). Unlike the CUDA build you do **not** set `CUDA_CC`: the HIP backend embeds a single generic `amdgcnspirv` image per kernel that the ROCm runtime finalizes for whichever AMD GPU is present at load time, so one build runs on any supported GPU. After editing a `.cu` kernel, regenerate the embedded images with `cd cuda && make wrappers BACKEND=hip` (this step needs `hipcc`; pass `HIPCC=/opt/rocm/bin/hipcc` if it is not on your `PATH`).
+
+The cgo flags default to a ROCm install at `/opt/rocm`; for a different prefix, set `CGO_CFLAGS`/`CGO_LDFLAGS` to its `include` and `lib` directories. The HIP backend builds on Windows ROCm as well, provided the process loads a full ROCm runtime (the one providing `amd_comgr`) so the generic image can be finalized.
+
+Validated on gfx90a (CDNA2), gfx1100 (RDNA3) and gfx1201 (RDNA4).
+
+👉 *Check ROCm installation with: `rocminfo` (lists your GPU) and `hipcc --version`.*
+
+</details>
+
 <details><summary><b><i>Install Go</i></b></summary>
 
 * Download and install from [the Go website](https://go.dev/doc/install).

bench/gpus.txt

@@ -3,6 +3,7 @@
 4.194304e+06 3.497176955517716e+07 1.694520362670856e-14    "MX150"
 4.194304e+06 5.527840160261479e+07 1.6945213680693907e-14   "GTX 860M"
 4.194304e+06 7.260693991786541e+07 1.694520362670856e-14    "GTX 1050 (mobile)"
+4.194304e+06 7.620596372190422e+07 1.69452004508381e-14     "RTX A400"
 4.194304e+06 8.191301445972674e+07 1.694520362670856e-14    "GTX 1050Ti (mobile)"
 4.194304e+06 8.490111356520656e+07 1.6945203523231987e-14   "RTX 2050 (mobile)"
 4.194304e+06 9.754137844792007e+07 1.6945208085986004e-14   "GTX 1650 (mobile)"

cmd/mumax3/gpuinfo.go

@@ -0,0 +1,20 @@
+//go:build !hip
+
+package main
+
+import (
+	"fmt"
+
+	"github.com/mumax/3/cuda"
+)
+
+// gpuInfoLine returns the GPU description printed at startup. The CUDA build
+// reports the compute capability of the PTX selected for this device.
+func gpuInfoLine() string {
+	return fmt.Sprintf("GPU info: %s, using cc=%d PTX", cuda.GPUInfo, cuda.UseCC)
+}
+
+// goBuildTags are the build tags forwarded to "go run" when executing a .go
+// input script, so the script is compiled with the same backend as this
+// binary. The default (CUDA) build needs none.
+const goBuildTags = ""

cmd/mumax3/gpuinfo_hip.go

@@ -0,0 +1,21 @@
+//go:build hip
+
+package main
+
+import (
+	"fmt"
+
+	"github.com/mumax/3/cuda"
+)
+
+// gpuInfoLine returns the GPU description printed at startup. The HIP build
+// loads one generic amdgcnspirv (SPIR-V) image that the runtime finalizes for
+// this device; GPUInfo already names the device's gfx arch.
+func gpuInfoLine() string {
+	return fmt.Sprintf("GPU info: %s, using generic amdgcnspirv image", cuda.GPUInfo)
+}
+
+// goBuildTags are the build tags forwarded to "go run" when executing a .go
+// input script, so the script is compiled with the same backend as this
+// binary. The HIP build must select the hip backend.
+const goBuildTags = "hip"

cmd/mumax3/main.go

@@ -131,7 +131,11 @@ func runScript(fname string) {
 
 func runGoFile(fname string) {
 	// pass through flags
-	flags := []string{"run", fname}
+	flags := []string{"run"}
+	if goBuildTags != "" {
+		flags = append(flags, "-tags", goBuildTags)
+	}
+	flags = append(flags, fname)
 	flag.Visit(func(f *flag.Flag) {
 		if f.Name != "o" {
 			flags = append(flags, fmt.Sprintf("-%v=%v", f.Name, f.Value))
@@ -169,7 +173,7 @@ func printVersion() {
 	engine.LogOut(engine.UNAME)
 	engine.LogOut(fmt.Sprintf("commit hash: %s", commitHash))
 	engine.LogOut(getCPUInfo())
-	engine.LogOut(fmt.Sprintf("GPU info: %s, using cc=%d PTX", cuda.GPUInfo, cuda.UseCC))
+	engine.LogOut(gpuInfoLine())
 	osInfo := fmt.Sprintf("OS  info: %s, Hostname: %s", getOSInfo(), getHostname())
 	engine.LogOut(osInfo)
 	engine.LogOut(fmt.Sprintf("Timestamp: %s", time.Now().Format("2006-01-02 15:04:05")))

cuda/Makefile

@@ -1,15 +1,29 @@
-# Builds mumax3 cuda kernels and create GO wrappers for the compute capabilities listed in $CUDA_CC. 
-# If $CUDA_CC is not defined, then $CUDA_CC is set to "50".
+# Builds the mumax3 GPU kernels and generates the Go wrappers.
 #
-# The ${CUDA_HOME}/bin/nvcc compiler is used to compile the cuda kernels. If CUDA_HOME is not defined
-# it will look for an nvidia compiler in $PATH instead.
+# Two backends are supported, selected with BACKEND:
+#
+#   BACKEND=cuda (default): compile each .cu to PTX text with nvcc, one per
+#   compute capability in $CUDA_CC (default "50"). cuda2go embeds the PTX into
+#   <name>_wrapper.go keyed by compute capability. This is the upstream NVIDIA
+#   build.
+#
+#   BACKEND=hip: compile each .cu to ONE generic SPIR-V (amdgcnspirv) image
+#   with "hipcc --genco --offload-arch=amdgcnspirv". This is the faithful
+#   analog of the CUDA path's embedded PTX: amdgcnspirv is a forward-compatible
+#   virtual ISA that the ROCm runtime finalizes (JITs) at module load for
+#   whatever GPU is present, so a single embedded image runs on any supported
+#   gfx arch with no per-arch matrix and no rebuild for new GPUs. cuda2go embeds
+#   the image bytes (base64) into <name>_wrapper_hip.go and the runtime loader
+#   hands it straight to hipModuleLoadData. CUDA_CC does NOT drive the HIP build;
+#   it stays meaningful only for the CUDA/nvcc path.
 #
 # Examples:
 #
-#    make
-#    make CUDA_CC=70
-#    make CUDA_CC="50 52 53 60 61 62 70 72 75 80 86"
-#    make CUDA_HOME="/usr/local/cuda-12.6" CUDA_CC="50 52 53 60 61 62 70 72 75 80 86"
+#    make                                          # nvcc, CC 50
+#    make CUDA_CC="50 60 70 80 86 90"              # nvcc, several CCs
+#    make CUDA_HOME=/usr/local/cuda-12.6 CUDA_CC=90
+#    make BACKEND=hip                              # hipcc, single amdgcnspirv image
+#    make BACKEND=hip HIPCC=/opt/rocm/bin/hipcc
 #
 # Different CUDA versions support different compute capabilities, as shown in the list below. See https://stackoverflow.com/a/28933055.
 # CUDA SDK 10.0 support for compute capability       30 32 35 37 50 52 53 60 61 62 70 72 75
@@ -22,15 +36,52 @@
 
 SHELL = /bin/bash
 
+BACKEND ?= cuda
+
+CUDAFILES := $(wildcard *.cu)
+
+ifeq ($(BACKEND),hip)
+
+ifeq ($(HIPCC),)
+	HIPCC=hipcc
+endif
+
+# --offload-arch=amdgcnspirv emits one generic SPIR-V image per kernel that the
+# ROCm runtime finalizes for the present GPU at hipModuleLoadData time; no gfx
+# arch list is needed (CUDA_CC is a CUDA/nvcc-path concept only).
+HIP_OFFLOAD_ARCH = amdgcnspirv
+
+# -include hip/hip_runtime.h gives the kernels the built-in symbols (blockIdx,
+# threadIdx, __global__, atomicAdd, float3, ...) that nvcc injects implicitly
+# but hipcc requires to be included; the .cu sources stay CUDA-spelled and
+# unmodified.
+# -I. lets the CUDA-spelled #include <cuComplex.h> resolve to the local compat
+# shim that maps cuComplex names to hipComplex.
+# -Wno-bitwise-instead-of-logical: clang flags the intentional bitwise-or of
+# booleans in the topological-charge stencils that nvcc accepts; keep -Werror
+# for everything else.
+HIPCCFLAGS = -O3 -Wall -Werror -Wno-bitwise-instead-of-logical --genco -include hip/hip_runtime.h -I.
+
+WRAPPERS := $(CUDAFILES:.cu=_wrapper_hip.go)
+
+%_wrapper_hip.go: %.cu cuda2go
+	@ rm -f $(basename $<)*.co
+	@ echo $(HIPCC) $(HIPCCFLAGS) --offload-arch=$(HIP_OFFLOAD_ARCH) $< -o $(basename $<).co
+	@      $(HIPCC) $(HIPCCFLAGS) --offload-arch=$(HIP_OFFLOAD_ARCH) $< -o $(basename $<).co
+	@ ./cuda2go -backend=hip $< > /dev/null
+	@ gofmt -w -s -l $@ > /dev/null
+
+else
+
 # When CUDA_HOME is not an environment variable and is not set on the command line, use the nvcc compiler
 # from the PATH
 ifeq ($(CUDA_HOME),)
 	NVCC=nvcc
-else 
-	NVCC=${CUDA_HOME}/bin/nvcc 
+else
+	NVCC=${CUDA_HOME}/bin/nvcc
 endif
 
-# When CUDA_CC is not an environment variable and is not set on the command line, use compute capability 3.0
+# When CUDA_CC is not an environment variable and is not set on the command line, use compute capability 5.0
 ifeq ($(CUDA_CC),)
 	CUDA_CC = 50 # Lowest supported CC for mumax3.12
 endif
@@ -49,39 +100,43 @@ endif
 
 NVCCFLAGS = $(NVCC_COMPATIBILITY_FLAGS) -ccbin=$(NVCC_CCBIN) --compiler-options -Werror --compiler-options -Wall -Xptxas -O3 -ptx
 
-CUDAFILES := $(wildcard *.cu)
 WRAPPERS := $(CUDAFILES:.cu=_wrapper.go)
 
+%_wrapper.go: %.cu cuda2go
+	@ rm -f $(basename $<)*.ptx
+	@ for cc in $(CUDA_CC); do \
+		echo $(NVCC) $(NVCCFLAGS) -arch=compute_$$cc -code=sm_$$cc $< -o $(basename $<)_$$cc.ptx ;\
+		     $(NVCC) $(NVCCFLAGS) -arch=compute_$$cc -code=sm_$$cc $< -o $(basename $<)_$$cc.ptx ;\
+	done
+	@ ./cuda2go -backend=cuda $< > /dev/null
+	@ gofmt -w -s -l $@ > /dev/null
+	@ cp $(basename $<)_wrapper.go $(basename $<)_wrapper.go_linux_cuda${CUDA_VERSION}.tmp
+
+endif
+
 
 .PHONY: all wrappers clean realclean
 
 
 all: wrappers
+ifeq ($(BACKEND),hip)
+	@echo "Built generic amdgcnspirv (SPIR-V) HIP images"
+else
 	@echo "Built with CUDA version ${CUDA_VERSION}"
+endif
 	go install -v
 
 
 wrappers: $(WRAPPERS)
 
 
-%_wrapper.go: %.cu cuda2go
-	@ rm -f $(basename $<)*.ptx
-	@ for cc in $(CUDA_CC); do \
-		echo $(NVCC) $(NVCCFLAGS) -arch=compute_$$cc -code=sm_$$cc $< -o $(basename $<)_$$cc.ptx ;\
-		     $(NVCC) $(NVCCFLAGS) -arch=compute_$$cc -code=sm_$$cc $< -o $(basename $<)_$$cc.ptx ;\
-	done
-	@ ./cuda2go $< > /dev/null
-	@ gofmt -w -s -l $@ > /dev/null
-	@ cp $(basename $<)_wrapper.go $(basename $<)_wrapper.go_linux_cuda${CUDA_VERSION}.tmp
-
-
 cuda2go: cuda2go.go
 	go build $<
 
 
 clean:
-	rm -vf *.ptx *.tmp
+	rm -vf *.ptx *.co *.tmp
 
 
 realclean:
-	rm -vf *_wrapper.go *.ptx cuda2go
+	rm -vf *_wrapper.go *_wrapper_hip.go *.ptx *.co cuda2go

cuda/atomicf.h

@@ -2,9 +2,33 @@
 #define _ATOMICF_H_
 
 // Atomic max of abs value.
+// The accumulator and all candidates are non-negative (b = fabs(b) and the
+// destination is seeded non-negative), so the IEEE-754 bit pattern is monotonic
+// with the float value and an integer max over the bits is correct.
+//
+// On NVIDIA (the #else branch, upstream code) this is a single integer
+// atomicMax on the float reinterpreted as int. On AMD CDNA an integer atomicMax
+// can be silently dropped on coarse-grained memory, so the HIP branch uses an
+// atomicCAS loop instead, which is honored on every ROCm coherence mode and is
+// equivalent for non-negative inputs.
+#ifdef __HIP_PLATFORM_AMD__
+inline __device__ void atomicFmaxabs(float* a, float b){
+	b = fabs(b);
+	int  bbits = __float_as_int(b);
+	int* aint  = (int*)(a);
+	int  old   = *aint;
+	int  assumed;
+	do {
+		assumed = old;
+		if (__int_as_float(assumed) >= b) break;
+		old = atomicCAS(aint, assumed, bbits);
+	} while (assumed != old);
+}
+#else
 inline __device__ void atomicFmaxabs(float* a, float b){
 	b = fabs(b);
 	atomicMax((int*)(a), *((int*)(&b)));
 }
+#endif
 
 #endif

cuda/cellindices_wrapper.go

@@ -1,3 +1,5 @@
+//go:build !hip
+
 package cuda
 
 /*

cuda/copypadmul2_wrapper.go

@@ -1,3 +1,5 @@
+//go:build !hip
+
 package cuda
 
 /*