deepspeedai--deepspeed/op_builder/xpu/builder.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import time
import importlib

try:
    # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed
    # if successful this also means we're doing a local install and not JIT compile path
    from op_builder import __deepspeed__  # noqa: F401 # type: ignore
    from op_builder.builder import OpBuilder
except ImportError:
    from deepspeed.ops.op_builder.builder import OpBuilder


class SYCLOpBuilder(OpBuilder):

    def builder(self):
        from torch.utils.cpp_extension import SyclExtension
        include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())]
        print("sycl sources = {}".format(self.sources()))
        sycl_ext = SyclExtension(name=self.absolute_name(),
                                 sources=self.strip_empty_entries(self.sources()),
                                 include_dirs=include_dirs,
                                 extra_compile_args={
                                     'cxx': self.strip_empty_entries(self.cxx_args()),
                                 },
                                 extra_link_args=self.strip_empty_entries(self.fixed_aotflags()))
        return sycl_ext

    def version_dependent_macros(self):
        try:
            from op_builder.builder import TORCH_MAJOR, TORCH_MINOR
        except ImportError:
            from deepspeed.ops.op_builder.builder import TORCH_MAJOR, TORCH_MINOR
        # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456
        version_ge_1_1 = []
        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
            version_ge_1_1 = ['-DVERSION_GE_1_1']
        version_ge_1_3 = []
        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
            version_ge_1_3 = ['-DVERSION_GE_1_3']
        version_ge_1_5 = []
        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4):
            version_ge_1_5 = ['-DVERSION_GE_1_5']
        return version_ge_1_1 + version_ge_1_3 + version_ge_1_5

    def _sycl_env_paths(self):
        """Find the SYCL include and lib directories from the Python environment.

        When using PyTorch XPU wheels, libsycl.so and SYCL headers are
        installed into the Python environment (e.g. conda env).  The system
        ``icpx`` compiler ships its own (potentially newer) SYCL headers and
        runtime.  To avoid ABI mismatches we must compile and link against the
        *same* SYCL version that PyTorch was built with.

        Returns (include_dir, lib_dir) – either or both may be ``None`` when
        the paths do not exist.
        """
        import sys
        prefix = sys.prefix  # e.g. /home/user/miniforge3/envs/myenv
        inc = os.path.join(prefix, 'include')
        lib = os.path.join(prefix, 'lib')
        sycl_inc = inc if os.path.isdir(os.path.join(inc, 'sycl')) else None
        sycl_lib = lib if os.path.isfile(os.path.join(lib, 'libsycl.so')) else None
        return sycl_inc, sycl_lib

    def cxx_args(self):
        cxx_flags = [
            '-fsycl',
            '-fsycl-targets=spir64',
            '-g',
            '-gdwarf-4',
            '-O3',
            '-std=c++17',
            '-fPIC',
            '-DMKL_ILP64',
            '-fno-strict-aliasing',
        ]
        # Use SYCL headers from the Python environment so that compiled code
        # references symbols present in the *environment's* libsycl.so rather
        # than the (possibly newer) system oneAPI installation.
        sycl_inc, _ = self._sycl_env_paths()
        if sycl_inc:
            cxx_flags = [f'-isystem', sycl_inc] + cxx_flags
        if os.environ.get('USE_MKL_GEMM'):
            cxx_flags.append('-DUSE_MKL_GEMM')
        return cxx_flags

    def extra_ldflags(self):
        import torch
        torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), 'lib')
        flags = [
            '-fPIC',
            '-fsycl',
            '-fsycl-targets=spir64',
            '-Xs "-options -cl-intel-enable-auto-large-GRF-mode"',
            '-fsycl-max-parallel-link-jobs=8',
            '-Wl,-export-dynamic',
            f'-L{torch_lib_dir}',
            f'-Wl,-rpath,{torch_lib_dir}',
        ]
        # Link against the Python environment's libsycl.so to match the
        # headers we compiled against (see cxx_args).
        _, sycl_lib = self._sycl_env_paths()
        if sycl_lib:
            flags = [f'-L{sycl_lib}', f'-Wl,-rpath,{sycl_lib}'] + flags
        return flags

    def fixed_aotflags(self):
        return [
            '-fsycl', '-fsycl-targets=spir64', '-fsycl-max-parallel-link-jobs=8',
            '-Xs "-options -cl-intel-enable-auto-large-GRF-mode"'
        ]

    def load(self, verbose=True):
        from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name  # noqa: F401
        from deepspeed.accelerator import get_accelerator
        if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name:
            return importlib.import_module(self.absolute_name())
        else:
            return self.jit_load(verbose)

    def jit_load(self, verbose=True):
        if not self.is_compatible(verbose):
            raise RuntimeError(
                f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}"
            )
        from torch.utils.cpp_extension import verify_ninja_availability
        try:
            verify_ninja_availability()
        except RuntimeError as e:
            raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") from e

        self.jit_mode = True
        from torch.utils.cpp_extension import load

        start_build = time.time()
        # Recognize relative paths as absolute paths for jit load

        sources = [self.deepspeed_src_path(path) for path in self.sources()]
        extra_include_paths = [self.deepspeed_src_path(path) for path in self.include_paths()]

        # Set CXX to icpx (Intel oneAPI DPC++ compiler) so that .cpp/.dp.cpp
        # files containing SYCL code are compiled with the SYCL-aware compiler.
        # PyTorch's cpp_extension only routes .sycl files to icpx by default.
        saved_env = {}
        for var in ('CXX', 'LIBRARY_PATH', 'CPATH'):
            saved_env[var] = os.environ.get(var)
        os.environ['CXX'] = 'icpx'

        # Point icpx at the Python environment's SYCL headers and libraries so
        # the compiled extension uses the same SYCL ABI as PyTorch.
        sycl_inc, sycl_lib = self._sycl_env_paths()
        if sycl_lib:
            lib_path = os.environ.get('LIBRARY_PATH', '')
            os.environ['LIBRARY_PATH'] = f'{sycl_lib}:{lib_path}' if lib_path else sycl_lib
        if sycl_inc:
            cpath = os.environ.get('CPATH', '')
            os.environ['CPATH'] = f'{sycl_inc}:{cpath}' if cpath else sycl_inc

        try:
            op_module = load(name=self.name,
                             sources=self.strip_empty_entries(sources),
                             extra_include_paths=self.strip_empty_entries(extra_include_paths),
                             extra_cflags=self.strip_empty_entries(self.cxx_args()),
                             extra_ldflags=self.strip_empty_entries(self.extra_ldflags()),
                             verbose=verbose)
        finally:
            # Restore original environment
            for var, val in saved_env.items():
                if val is None:
                    os.environ.pop(var, None)
                else:
                    os.environ[var] = val

        build_duration = time.time() - start_build
        if verbose:
            print(f"Time to load {self.name} op: {build_duration} seconds")
        return op_module