Add support for AMD GPUs via the ProcessAcceleratorROCm

Author: fwyzard

Configuration/StandardSequences/python/Accelerators_cff.py

@@ -4,3 +4,4 @@
 # used in production
 
 from HeterogeneousCore.CUDACore.ProcessAcceleratorCUDA_cfi import ProcessAcceleratorCUDA
+from HeterogeneousCore.ROCmCore.ProcessAcceleratorROCm_cfi import ProcessAcceleratorROCm

HeterogeneousCore/AlpakaCore/python/ProcessAcceleratorAlpaka.py

@@ -2,10 +2,12 @@
 
 class ModuleTypeResolverAlpaka:
     def __init__(self, accelerators, backend):
-        # first element is used as the default is nothing is set
+        # first element is used as the default if nothing is set
         self._valid_backends = []
         if "gpu-nvidia" in accelerators:
             self._valid_backends.append("cuda_async")
+        if "gpu-amd" in accelerators:
+            self._valid_backends.append("rocm_async")
         if "cpu" in accelerators:
             self._valid_backends.append("serial_sync")
         if len(self._valid_backends) == 0:
@@ -60,11 +62,15 @@ def apply(self, process, accelerators):
         if not hasattr(process, "AlpakaServiceCudaAsync"):
             from HeterogeneousCore.AlpakaServices.AlpakaServiceCudaAsync_cfi import AlpakaServiceCudaAsync
             process.add_(AlpakaServiceCudaAsync)
+        if not hasattr(process, "AlpakaServiceROCmAsync"):
+            from HeterogeneousCore.AlpakaServices.AlpakaServiceROCmAsync_cfi import AlpakaServiceROCmAsync
+            process.add_(AlpakaServiceROCmAsync)
 
         if not hasattr(process.MessageLogger, "AlpakaService"):
             process.MessageLogger.AlpakaService = cms.untracked.PSet()
 
         process.AlpakaServiceSerialSync.enabled = "cpu" in accelerators
         process.AlpakaServiceCudaAsync.enabled = "gpu-nvidia" in accelerators
+        process.AlpakaServiceROCmAsync.enabled = "gpu-amd" in accelerators
 
 cms.specialImportRegistry.registerSpecialImportForType(ProcessAcceleratorAlpaka, "from HeterogeneousCore.AlpakaCore.ProcessAcceleratorAlpaka import ProcessAcceleratorAlpaka")

HeterogeneousCore/ROCmCore/python/ProcessAcceleratorROCm.py

@@ -0,0 +1,30 @@
+import FWCore.ParameterSet.Config as cms
+
+import os
+
+class ProcessAcceleratorROCm(cms.ProcessAccelerator):
+    def __init__(self):
+        super(ProcessAcceleratorROCm,self).__init__()
+        self._label = "gpu-amd"
+    def labels(self):
+        return [self._label]
+    def enabledLabels(self):
+        enabled = (os.system("rocmIsEnabled") == 0)
+        if enabled:
+            return self.labels()
+        else:
+            return []
+    def apply(self, process, accelerators):
+        if not hasattr(process, "ROCmService"):
+            from HeterogeneousCore.ROCmServices.ROCmService_cfi import ROCmService
+            process.add_(ROCmService)
+
+        if not hasattr(process.MessageLogger, "ROCmService"):
+            process.MessageLogger.ROCmService = cms.untracked.PSet()
+
+        if self._label in accelerators:
+            process.ROCmService.enabled = True
+        else:
+            process.ROCmService.enabled = False
+            
+cms.specialImportRegistry.registerSpecialImportForType(ProcessAcceleratorROCm, "from HeterogeneousCore.ROCmCore.ProcessAcceleratorROCm import ProcessAcceleratorROCm")

HeterogeneousCore/ROCmCore/python/ProcessAcceleratorROCm_cfi.py

@@ -0,0 +1,3 @@
+from HeterogeneousCore.ROCmCore.ProcessAcceleratorROCm import ProcessAcceleratorROCm as _ProcessAcceleratorROCm
+
+ProcessAcceleratorROCm = _ProcessAcceleratorROCm()