setup.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

tnn_utils = CUDAExtension(
    name="tt_embeddings",
    sources=[
        "tt_embeddings.cpp",
        "tt_embeddings_cuda.cu",
    ],
    extra_compile_args={
        "cxx": [
            "-O3",
            "-g",
            "-DUSE_MKL",
            "-m64",
            "-mfma",
            "-masm=intel",
        ],
        "nvcc": [
            "-O3",
            "-g",
            "--expt-relaxed-constexpr",
            "-D__CUDA_NO_HALF_OPERATORS__",
            "-I/usr/lib/cub-1.8.0",
            '-gencode=arch=compute_70,code="sm_70"',
        ],
    },
)
setup(
    name="tt_embeddings",
    description="tt_embeddings",
    packages=find_packages(),
    ext_modules=[tnn_utils],
    cmdclass={"build_ext": BuildExtension},
)