Repository: microsoft/DeepSpeed
File: op_builder/builder.py (Python)

Code before:

    @staticmethod
    def get_rocm_gpu_arch():
        if OpBuilder._rocm_gpu_arch:
            return OpBuilder._rocm_gpu_arch
        rocm_info = Path("/opt/rocm/bin/rocminfo")
        if (not rocm_info.is_file()):
            rocm_info = Path("rocminfo")
        rocm_gpu_arch_cmd = str(rocm_info) + " | grep -o -m 1 'gfx.*'"
        try:
            result = subprocess.check_output(rocm_gpu_arch_cmd, shell=True)
            rocm_gpu_arch = result.decode('utf-8').strip()
        except subprocess.CalledProcessError:
            rocm_gpu_arch = ""
        OpBuilder._rocm_gpu_arch = rocm_gpu_arch
        return OpBuilder._rocm_gpu_arch

    @staticmethod
    def get_rocm_wavefront_size():
        if OpBuilder._rocm_wavefront_size:
            return OpBuilder._rocm_wavefront_size

        rocm_info = Path("/opt/rocm/bin/rocminfo")
        if (not rocm_info.is_file()):
            rocm_info = Path("rocminfo")
        rocm_wavefront_size_cmd = str(rocm_info) + " | grep -Eo -m1 'Wavefront Size:[[:space:]]+[0-9]+' | grep -Eo '[0-9]+'"
        try:
            result = subprocess.check_output(rocm_wavefront_size_cmd, shell=True)
            rocm_wavefront_size = result.decode('utf-8').strip()
        except subprocess.CalledProcessError:
            rocm_wavefront_size = "32"
        OpBuilder._rocm_wavefront_size = rocm_wavefront_size
        return OpBuilder._rocm_wavefront_size

    def include_paths(self):
        '''
        Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed)
        '''
        return []

    def nvcc_args(self):
        '''
        Returns optional list of compiler flags to forward to nvcc when building CUDA sources
        '''
        return []

    def cxx_args(self):
        '''
        Returns optional list of compiler flags to forward to the build
        '''
        return []

Review comment (suggestion):
This should become available in `torch.cuda.get_device_properties` with https://github.com/pytorch/pytorch/pull/128449 

Code after:

    @staticmethod
    def get_rocm_gpu_arch():
        if OpBuilder._rocm_gpu_arch:
            return OpBuilder._rocm_gpu_arch
        rocm_info = Path("/opt/rocm/bin/rocminfo")
        if (not rocm_info.is_file()):
            rocm_info = Path("rocminfo")
        rocm_gpu_arch_cmd = str(rocm_info) + " | grep -o -m 1 'gfx.*'"
        try:
            result = subprocess.check_output(rocm_gpu_arch_cmd, shell=True)
            rocm_gpu_arch = result.decode('utf-8').strip()
        except subprocess.CalledProcessError:
            rocm_gpu_arch = ""
        OpBuilder._rocm_gpu_arch = rocm_gpu_arch
        return OpBuilder._rocm_gpu_arch

    @staticmethod
    def get_rocm_wavefront_size():
        if OpBuilder._rocm_wavefront_size:
            return OpBuilder._rocm_wavefront_size

        rocm_info = Path("/opt/rocm/bin/rocminfo")
        if (not rocm_info.is_file()):
            rocm_info = Path("rocminfo")
        rocm_wavefront_size_cmd = str(
            rocm_info) + " | grep -Eo -m1 'Wavefront Size:[[:space:]]+[0-9]+' | grep -Eo '[0-9]+'"
        try:
            result = subprocess.check_output(rocm_wavefront_size_cmd, shell=True)
            rocm_wavefront_size = result.decode('utf-8').strip()
        except subprocess.CalledProcessError:
            rocm_wavefront_size = "32"
        OpBuilder._rocm_wavefront_size = rocm_wavefront_size
        return OpBuilder._rocm_wavefront_size

    def include_paths(self):
        '''
        Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed)
        '''
        return []

    def nvcc_args(self):
        '''
        Returns optional list of compiler flags to forward to nvcc when building CUDA sources
        '''
        return []

    def cxx_args(self):
        '''
        Returns optional list of compiler flags to forward to the build
        '''