File indexing completed on 2023-03-17 11:05:41
0001 import FWCore.ParameterSet.Config as cms
0002
0003 import os
0004
0005 from HeterogeneousCore.Common.PlatformStatus import PlatformStatus
0006
0007 class ModuleTypeResolverAlpaka:
0008 def __init__(self, accelerators, backend):
0009
0010 self._valid_backends = []
0011 if "gpu-nvidia" in accelerators:
0012 self._valid_backends.append("cuda_async")
0013 if "gpu-amd" in accelerators:
0014 self._valid_backends.append("rocm_async")
0015 if "cpu" in accelerators:
0016 self._valid_backends.append("serial_sync")
0017 if len(self._valid_backends) == 0:
0018 raise cms.EDMException(cms.edm.errors.UnavailableAccelerator, "ModuleTypeResolverAlpaka had no backends available because of the combination of the job configuration and accelerator availability of on the machine. The job sees {} accelerators.".format(", ".join(accelerators)))
0019 if backend is not None:
0020 if not backend in self._valid_backends:
0021 raise cms.EDMException(cms.edm.errors.UnavailableAccelerator, "The ProcessAcceleratorAlpaka was configured to use {} backend, but that backend is not available because of the combination of the job configuration and accelerator availability on the machine. The job was configured to use {} accelerators, which translates to {} Alpaka backends.".format(
0022 backend, ", ".join(accelerators), ", ".join(self._valid_backends)))
0023 if backend != self._valid_backends[0]:
0024 self._valid_backends.remove(backend)
0025 self._valid_backends.insert(0, backend)
0026
0027 def plugin(self):
0028 return "ModuleTypeResolverAlpaka"
0029
0030 def setModuleVariant(self, module):
0031 if module.type_().endswith("@alpaka"):
0032 defaultBackend = self._valid_backends[0]
0033 if hasattr(module, "alpaka"):
0034 if hasattr(module.alpaka, "backend"):
0035 if module.alpaka.backend == "":
0036 module.alpaka.backend = defaultBackend
0037 elif module.alpaka.backend.value() not in self._valid_backends:
0038 raise cms.EDMException(cms.edm.errors.UnavailableAccelerator, "Module {} has the Alpaka backend set explicitly, but its accelerator is not available for the job because of the combination of the job configuration and accelerator availability on the machine. The following Alpaka backends are available for the job {}.".format(module.label_(), ", ".join(self._valid_backends)))
0039 else:
0040 module.alpaka.backend = cms.untracked.string(defaultBackend)
0041 else:
0042 module.alpaka = cms.untracked.PSet(
0043 backend = cms.untracked.string(defaultBackend)
0044 )
0045
0046 class ProcessAcceleratorAlpaka(cms.ProcessAccelerator):
0047 """ProcessAcceleratorAlpaka itself does not define or inspect
0048 availability of any accelerator devices. It merely sets up
0049 necessary Alpaka infrastructure based on the availability of
0050 accelerators that the concrete ProcessAccelerators (like
0051 ProcessAcceleratorCUDA) define.
0052 """
0053 def __init__(self):
0054 super(ProcessAcceleratorAlpaka, self).__init__()
0055 self._backend = None
0056
0057
0058 def setBackend(self, backend):
0059 self._backend = backend
0060
0061
0062 def moduleTypeResolver(self, accelerators):
0063 return ModuleTypeResolverAlpaka(accelerators, self._backend)
0064
0065 def apply(self, process, accelerators):
0066
0067 if not hasattr(process.MessageLogger, "AlpakaService"):
0068 process.MessageLogger.AlpakaService = cms.untracked.PSet()
0069
0070
0071 try:
0072 if not "cpu" in accelerators:
0073 raise False
0074 from HeterogeneousCore.AlpakaServices.AlpakaServiceSerialSync_cfi import AlpakaServiceSerialSync
0075 except:
0076
0077 if hasattr(process, "AlpakaServiceSerialSync"):
0078 del process.AlpakaServiceSerialSync
0079 else:
0080
0081 if not hasattr(process, "AlpakaServiceSerialSync"):
0082 process.add_(AlpakaServiceSerialSync)
0083
0084
0085 try:
0086 if not "gpu-nvidia" in accelerators:
0087 raise False
0088 from HeterogeneousCore.AlpakaServices.AlpakaServiceCudaAsync_cfi import AlpakaServiceCudaAsync
0089 except:
0090
0091 if hasattr(process, "AlpakaServiceCudaAsync"):
0092 del process.AlpakaServiceCudaAsync
0093 else:
0094
0095 if not hasattr(process, "AlpakaServiceCudaAsync"):
0096 process.add_(AlpakaServiceCudaAsync)
0097
0098
0099 try:
0100 if not "gpu-amd" in accelerators:
0101 raise False
0102 from HeterogeneousCore.AlpakaServices.AlpakaServiceROCmAsync_cfi import AlpakaServiceROCmAsync
0103 except:
0104
0105 if hasattr(process, "AlpakaServiceROCmAsync"):
0106 del process.AlpakaServiceROCmAsync
0107 else:
0108
0109 if not hasattr(process, "AlpakaServiceROCmAsync"):
0110 process.add_(AlpakaServiceROCmAsync)
0111
0112
0113
0114 cms.specialImportRegistry.registerSpecialImportForType(ProcessAcceleratorAlpaka, "from HeterogeneousCore.AlpakaCore.ProcessAcceleratorAlpaka import ProcessAcceleratorAlpaka")