defs.bzl

# Copyright (c) Meta Platforms, Inc. and affiliates.
# Copyright 2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

def default_compiler_flags():
    return [
        "-pedantic",
        "-Wall",
        "-DFBGEMM_FBCODE",  # Used for enabling fbcode-specific code paths
    ] + select({
        "DEFAULT": [],
        "ovr_config//compiler:clang": [
            "-Werror",
            "-Wextra",
            "-Wno-gnu-zero-variadic-macro-arguments",
            "-Wno-c99-extensions",
            "-Wimplicit-fallthrough",
            "-Wignored-qualifiers",
            "-Wno-vla",
        ],
        "ovr_config//compiler:gcc": [
            "-Werror",
            "-Wextra",
            # Buck exposes Folly as non-system headers, where its intentional
            # `__int128` use triggers the `-pedantic` diagnostic.
            "-Wno-error=pedantic",
            "-fno-trapping-math",
        ]
        + select({
            "DEFAULT": [],
            "ovr_config//cpu:arm64": [
                "-ftree-vectorize",
            ],
        }),
    })

def get_fbgemm_base_srcs():
    return [
        "src/GenerateI8Depthwise.cc",
        "src/JitPerfMap.cc",
        "src/RefImplementations.cc",
        "src/Utils.cc",
    ]

def get_fbgemm_generic_srcs(with_base = False):
    sources = [
        "src/EmbeddingSpMDM.cc",
        "src/EmbeddingSpMDMNBit.cc",
        "src/ExecuteKernel.cc",
        "src/ExecuteKernelU8S8.cc",
        "src/Fbgemm.cc",
        "src/FbgemmBfloat16Convert.cc",
        "src/FbgemmConv.cc",
        "src/FbgemmFPCommon.cc",
        "src/FbgemmFP16.cc",
        "src/fp32/FbgemmFP32.cc",
        "src/FbgemmFloat16Convert.cc",
        "src/FbgemmI64.cc",
        "src/FbgemmSparseDense.cc",
        "src/FbgemmI8Spmdm.cc",
        "src/FbgemmPackMatrixB.cc",
        "src/GenerateKernelDirectConvU8S8S32ACC32.cc",
        "src/GenerateKernel.cc",
        "src/GenerateKernelU8S8S32ACC16.cc",
        "src/GenerateKernelU8S8S32ACC16Avx512.cc",  # Acc16 AVX512 JIT code gen
        "src/GenerateKernelU8S8S32ACC16Avx512VNNI.cc",
        "src/GenerateKernelU8S8S32ACC32.cc",
        "src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc",
        "src/GroupwiseConv.cc",
        "src/GroupwiseConvAcc32Avx2.cc",
        "src/GroupwiseConvAcc32Avx512.cc",
        "src/PackAMatrix.cc",
        "src/PackAWithIm2Col.cc",
        "src/PackAWithQuantRowOffset.cc",
        "src/PackAWithRowOffset.cc",
        "src/PackBMatrix.cc",
        "src/PackMatrix.cc",
        "src/PackWeightMatrixForGConv.cc",
        "src/PackWeightsForConv.cc",
        "src/PackWeightsForDirectConv.cc",
        "src/QuantUtils.cc",
        "src/RowWiseSparseAdagradFused.cc",
        "src/SparseAdagrad.cc",
        "src/spmmUtils.cc",
        "src/TransposeUtils.cc",
    ] + (get_fbgemm_base_srcs() if with_base else [])

    return sources

def get_fbgemm_public_headers():
    return [
        "include/fbgemm/Assert.h",
        "include/fbgemm/ConvUtils.h",
        "include/fbgemm/Fbgemm.h",
        "include/fbgemm/FbgemmBuild.h",
        "include/fbgemm/FbgemmConvert.h",
        "include/fbgemm/FbgemmEmbedding.h",
        "include/fbgemm/detail/EmbeddingSpMDMFP8.h",
        "include/fbgemm/FbgemmFP16.h",
        "include/fbgemm/FbgemmFP32.h",
        "include/fbgemm/FbgemmFPCommon.h",
        "include/fbgemm/FbgemmI64.h",
        "include/fbgemm/FbgemmI8Depthwise.h",
        "include/fbgemm/FbgemmI8DirectconvAvx2.h",
        "include/fbgemm/FbgemmI8Spmdm.h",
        "include/fbgemm/FbgemmPackMatrixB.h",
        "include/fbgemm/FbgemmSparse.h",
        "include/fbgemm/FloatConversion.h",
        "include/fbgemm/OutputProcessing-inl.h",
        "include/fbgemm/PackingTraits-inl.h",
        "include/fbgemm/QuantUtils.h",
        "include/fbgemm/QuantUtilsAvx2.h",
        "include/fbgemm/QuantUtilsAvx512.h",
        "include/fbgemm/QuantUtilsNeon.h",
        "include/fbgemm/spmmUtils.h",
        "include/fbgemm/spmmUtilsAvx2.h",
        "include/fbgemm/SimdUtils.h",
        "include/fbgemm/Utils.h",
        "include/fbgemm/UtilsAvx2.h",
        "include/fbgemm/Types.h",
    ]

def get_fbgemm_avx2_srcs(buck = False):
    # These 3 depthwise files are also compiled in fbgemm_sve.  On ARM,
    # exclude them from fbgemm_avx2 to avoid duplicate symbols when
    # downstream targets pull in both fbgemm_avx2 and fbgemm_sve (the
    # latter is selected by the main fbgemm target on arm64).
    depthwise_srcs = [
        "src/FbgemmI8Depthwise3D.cc",
        "src/FbgemmI8Depthwise.cc",
        "src/PackDepthwiseConvMatrix.cc",
    ]

    common_srcs = [
        # All the source files that either use avx2 instructions statically
        "src/EmbeddingSpMDMAvx2.cc",
        "src/FbgemmBfloat16ConvertAvx2.cc",
        "src/FbgemmFloat16ConvertAvx2.cc",
        "src/FbgemmI8DepthwisePerChannelQuantAvx2.cc",
        "src/FbgemmSparseDenseAvx2.cc",
        "src/FbgemmSparseDenseInt8Avx2.cc",
        "src/OptimizedKernelsAvx2.cc",
        "src/QuantUtilsAvx2.cc",
        "src/spmmUtilsAvx2.cc",
        "src/UtilsAvx2.cc",
    ]

    if buck:
        return common_srcs + select({
            "DEFAULT": depthwise_srcs,
            "ovr_config//cpu:arm64": [],  # Provided by fbgemm_sve on ARM
        })
    return common_srcs + depthwise_srcs

def get_fbgemm_inline_avx2_srcs(msvc = False, buck = False):
    intrinsics_srcs = ["src/FbgemmFP16UKernelsIntrinsicAvx2.cc"]

    # FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
    asm_srcs = [
        "src/fp32/FbgemmFP32UKernelsAvx2.cc",
        "src/FbgemmFP16UKernelsAvx2.cc",
    ]
    if buck:
        return select({
            "DEFAULT": asm_srcs,
            "ovr_config//compiler:cl": intrinsics_srcs,
            "ovr_config//cpu:arm64": intrinsics_srcs,
        })
    return asm_srcs if not msvc else intrinsics_srcs

def get_fbgemm_avx512_srcs():
    return [
        # All the source files that use avx512 instructions statically
        "src/FbgemmBfloat16ConvertAvx512.cc",
        "src/EmbeddingSpMDMAvx512.cc",
        "src/FbgemmFloat16ConvertAvx512.cc",
        "src/FbgemmSparseDenseAvx512.cc",
        "src/FbgemmSparseDenseInt8Avx512.cc",
        "src/FbgemmSparseDenseVectorInt8Avx512.cc",
        "src/QuantUtilsAvx512.cc",
        "src/UtilsAvx512.cc",
    ]

def get_fbgemm_inline_avx512_srcs(msvc = False, buck = False):
    intrinsics_srcs = [
        "src/FbgemmFP16UKernelsIntrinsicAvx512.cc",
        "src/FbgemmFP16UKernelsIntrinsicAvx512_256.cc",
    ]
    asm_srcs = [
        "src/FbgemmFP16UKernelsAvx512.cc",
        "src/FbgemmFP16UKernelsAvx512_256.cc",
        "src/fp32/FbgemmFP32UKernelsAvx512.cc",
        "src/fp32/FbgemmFP32UKernelsAvx512_256.cc",
    ]
    if buck:
        return select({
            "DEFAULT": asm_srcs,
            "ovr_config//compiler:cl": intrinsics_srcs,
            "ovr_config//cpu:arm64": intrinsics_srcs,
        })
    return asm_srcs if not msvc else intrinsics_srcs

def get_fbgemm_inline_sve_srcs(msvc = False, buck = False):
    srcs = [
        "src/FbgemmFP16UKernelsSve128.cc",
        "src/UtilsSve.cc",
        "src/FbgemmFloat16ConvertSVE.cc",
        "src/PackDepthwiseConvMatrix.cc",
        "src/FbgemmI8Depthwise3D.cc",
        "src/FbgemmI8Depthwise.cc",
    ]

    if buck:
        # FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
        srcs += select({
            "DEFAULT": [],
            "ovr_config//cpu:arm64": [
                "src/FbgemmFloat16ConvertSVE.cc",
            ],
        })
        return select({
            "DEFAULT": srcs,
            "ovr_config//compiler:cl": srcs,
            "ovr_config//cpu:arm64": srcs,
        })

    else:
        return srcs

def get_fbgemm_inline_neon_srcs(msvc = False, buck = False):
    intrinsics_srcs = [
        "src/UtilsNeon.cc",
        "src/FbgemmFloat16ConvertNEON.cc",
    ]

    # FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
    asm_srcs = [
        "src/UtilsNeon.cc",
        "src/FbgemmFloat16ConvertNEON.cc",
        "src/KleidiAIFP16UKernelsNeon.cc",
        "src/fp32/KleidiAIFP32UKernelsNeon.cc",
        "src/QuantUtilsNeon.cc",
    ]
    if buck:
        return select({
            "DEFAULT": asm_srcs,
            "ovr_config//compiler:cl": intrinsics_srcs,
            "ovr_config//cpu:arm64": asm_srcs,
        })
    return asm_srcs if not msvc else intrinsics_srcs

def get_fbgemm_autovec_srcs():
    return [
        "src/EmbeddingSpMDMAutovec.cc",
        "src/EmbeddingStatsTracker.cc",
    ]

def get_fbgemm_tests(skip_tests = []):
    return native.glob(["test/*Test.cc"], exclude = skip_tests)

def select_fbgemm_codegen_mode(inference, training):
    return select({
        "//deeplearning/fbgemm/constraints:codegen_mode[inference]": inference,
        "//deeplearning/fbgemm/constraints:codegen_mode[training]": training,
        "DEFAULT": select({
            "//deeplearning/fbgemm/constraints:fbgemm_codegen_inference_mode_enabled": inference,
            "//deeplearning/fbgemm/constraints:fbgemm_codegen_inference_mode_enabled_capitalized": inference,
            "//deeplearning/fbgemm/constraints:fbgemm_codegen_inference_mode_enabled_numeric": inference,
            "DEFAULT": training,
        }),
    })