HeterogeneousCore/AlpakaCore/python/ProcessAcceleratorAlpaka.py - master - githubmirror/cmssw

HeterogeneousCore/AlpakaCore/python/ProcessAcceleratorAlpaka.py
130 строк7 KB
Andrea Bocci
Include the ProcessAcceleratorAlpaka configuration in the python dumps
30 апр 2026, 18:43
30 апр 2026, 18:43367d7da
import FWCore.ParameterSet.Config as cms

import os

from HeterogeneousCore.Common.PlatformStatus import PlatformStatus

class ModuleTypeResolverAlpaka:
    def __init__(self, accelerators, backend, synchronize):
        # 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:
            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)))
        if backend is not None:
            if not backend in self._valid_backends:
                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(
                    backend, ", ".join(accelerators), ", ".join(self._valid_backends)))
            if backend != self._valid_backends[0]:
                self._valid_backends.remove(backend)
                self._valid_backends.insert(0, backend)
        self._synchronize = synchronize

    def plugin(self):
        return "ModuleTypeResolverAlpaka"

    def setModuleVariant(self, module):
        if module.type_().endswith("@alpaka"):
            defaultBackend = self._valid_backends[0]
            if hasattr(module, "alpaka"):
                # Ensure the untrackedness already here, because the
                # C++ ModuleTypeResolverAlpaka relies on the
                # untrackedness (before the configuration validation)
                if module.alpaka.isTracked():
                    raise cms.EDMException(cms.edm.errors.Configuration, "The 'alpaka' PSet in module '{}' is tracked, but it should be untracked".format(module.label()))
                if hasattr(module.alpaka, "backend"):
                    if module.alpaka.backend == "":
                        module.alpaka.backend = defaultBackend
                    elif module.alpaka.backend.value() not in self._valid_backends:
                        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)))
                else:
                    module.alpaka.backend = cms.untracked.string(defaultBackend)
            else:
                module.alpaka = cms.untracked.PSet(
                    backend = cms.untracked.string(defaultBackend)
                )
            isDefaultValue = lambda v: \
                isinstance(v, type(cms.optional.untracked.bool)) \
                and not v.isTracked() \
                and v.isCompatibleCMSType(cms.bool)
            if not hasattr(module.alpaka, "synchronize") or isDefaultValue(module.alpaka.synchronize):
                module.alpaka.synchronize = cms.untracked.bool(self._synchronize)

class ProcessAcceleratorAlpaka(cms.ProcessAccelerator):
    """ProcessAcceleratorAlpaka itself does not define or inspect
    availability of any accelerator devices. It merely sets up
    necessary Alpaka infrastructure based on the availability of
    accelerators that the concrete ProcessAccelerators (like
    ProcessAcceleratorCUDA) define.
    """
    def __init__(self, backend = None, synchronize = False):
        super(ProcessAcceleratorAlpaka, self).__init__()
        self._backend = backend
        self._synchronize = synchronize

    # User-facing interface
    def setBackend(self, backend):
        self._backend = backend

    def setSynchronize(self, synchronize):
        self._synchronize = synchronize

    # Include the ProcessAcceleratorAlpaka configuration in the python dumps
    def dumpPythonImpl(self, options) -> str:
        args = []
        if self._backend is not None:
            args.append(options.indentation() + f"backend = '{self._backend}'")
        if self._synchronize != False:
            args.append(options.indentation() + f"synchronize = {self._synchronize}")
        return ',\n'.join(args)

    # Framework-facing interface
    def moduleTypeResolver(self, accelerators):
        return ModuleTypeResolverAlpaka(accelerators, self._backend, self._synchronize)

    def apply(self, process, accelerators):
        # Propagate the AlpakaService messages through the MessageLogger
        if not hasattr(process.MessageLogger, "AlpakaService"):
            process.MessageLogger.AlpakaService = cms.untracked.PSet()

        # The CPU backend is effectively always available, ensure the AlpakaServiceSerialSync is loaded
        if not hasattr(process, "AlpakaServiceSerialSync"):
            from HeterogeneousCore.AlpakaServices.AlpakaServiceSerialSync_cfi import AlpakaServiceSerialSync
            process.add_(AlpakaServiceSerialSync)

        # Check if CUDA is available, and if the system has at least one usable NVIDIA GPU
        try:
            if not "gpu-nvidia" in accelerators:
                raise False
            from HeterogeneousCore.AlpakaServices.AlpakaServiceCudaAsync_cfi import AlpakaServiceCudaAsync
        except:
            # CUDA is not available, do not load the AlpakaServiceCudaAsync
            if hasattr(process, "AlpakaServiceCudaAsync"):
                del process.AlpakaServiceCudaAsync
        else:
            # CUDA is available, ensure the AlpakaServiceCudaAsync is loaded
            if not hasattr(process, "AlpakaServiceCudaAsync"):
                process.add_(AlpakaServiceCudaAsync)

        # Check if ROCm is available, and if the system has at least one usable AMD GPU
        try:
            if not "gpu-amd" in accelerators:
                raise False
            from HeterogeneousCore.AlpakaServices.AlpakaServiceROCmAsync_cfi import AlpakaServiceROCmAsync
        except:
            # ROCm is not available, do not load the AlpakaServiceROCmAsync
            if hasattr(process, "AlpakaServiceROCmAsync"):
                del process.AlpakaServiceROCmAsync
        else:
            # ROCm is available, ensure the AlpakaServiceROCmAsync is loaded
            if not hasattr(process, "AlpakaServiceROCmAsync"):
                process.add_(AlpakaServiceROCmAsync)


# Ensure this module is kept in the configuration when dumping it
cms.specialImportRegistry.registerSpecialImportForType(ProcessAcceleratorAlpaka, "from HeterogeneousCore.AlpakaCore.ProcessAcceleratorAlpaka import ProcessAcceleratorAlpaka")