"docs/vscode:/vscode.git/clone" did not exist on "a0939977a3b3c34c925c565c3fd3dcbe5d09e23c"
build.py 2.23 KB
Newer Older
Hang Zhang's avatar
init  
Hang Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree 
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import os
import torch
Hang Zhang's avatar
Hang Zhang committed
13
import platform
Hang Zhang's avatar
backend  
Hang Zhang committed
14
import subprocess
Hang Zhang's avatar
init  
Hang Zhang committed
15
16
from torch.utils.ffi import create_extension

Hang Zhang's avatar
Hang Zhang committed
17
18
19
lib_path = os.path.join(os.path.dirname(torch.__file__), 'lib')
this_file = os.path.dirname(os.path.realpath(__file__))

Hang Zhang's avatar
backend  
Hang Zhang committed
20
# build kernel library
Hang Zhang's avatar
Hang Zhang committed
21
22
os.environ['TORCH_BUILD_DIR'] = lib_path
if platform.system() == 'Darwin':
Hang Zhang's avatar
indent  
Hang Zhang committed
23
24
25
    os.environ['TH_LIBRARIES'] = os.path.join(lib_path,'libTH.1.dylib')
    os.environ['THC_LIBRARIES'] = os.path.join(lib_path,'libTHC.1.dylib')
    ENCODING_LIB = os.path.join(lib_path, 'libENCODING.dylib')
Hang Zhang's avatar
Hang Zhang committed
26
else:
Hang Zhang's avatar
indent  
Hang Zhang committed
27
28
29
    os.environ['TH_LIBRARIES'] = os.path.join(lib_path,'libTH.so.1')
    os.environ['THC_LIBRARIES'] = os.path.join(lib_path,'libTHC.so.1')
    ENCODING_LIB = os.path.join(lib_path, 'libENCODING.so')
Hang Zhang's avatar
Hang Zhang committed
30

31
32
33
clean_cmd = ['bash', 'clean.sh']
subprocess.check_call(clean_cmd)

Hang Zhang's avatar
backend  
Hang Zhang committed
34
build_all_cmd = ['bash', 'encoding/make.sh']
35
subprocess.check_call(build_all_cmd, env=dict(os.environ))
Hang Zhang's avatar
backend  
Hang Zhang committed
36
37
38
39
40
41

sources = ['encoding/src/encoding_lib.cpp']
headers = ['encoding/src/encoding_lib.h']
defines = [('WITH_CUDA', None)]
with_cuda = True 

Hang Zhang's avatar
Hang Zhang committed
42
include_path = [os.path.join(lib_path, 'include'),
Hang Zhang's avatar
indent  
Hang Zhang committed
43
44
45
                                os.path.join(os.environ['HOME'],'pytorch/torch/lib/THC'), 
                                os.path.join(lib_path,'include/ENCODING'), 
                                os.path.join(this_file,'encoding/src/')]
Hang Zhang's avatar
init  
Hang Zhang committed
46

Hang Zhang's avatar
Hang Zhang committed
47
def make_relative_rpath(path):
Hang Zhang's avatar
indent  
Hang Zhang committed
48
49
50
51
    if platform.system() == 'Darwin':
        return '-Wl,-rpath,' + path
    else:
        return '-Wl,-rpath,' + path
Hang Zhang's avatar
tested  
Hang Zhang committed
52

Hang Zhang's avatar
init  
Hang Zhang committed
53
ffi = create_extension(
Hang Zhang's avatar
indent  
Hang Zhang committed
54
55
56
57
58
59
60
61
62
63
64
65
    'encoding._ext.encoding_lib',
    package=True,
    headers=headers,
    sources=sources,
    define_macros=defines,
    relative_to=__file__,
    with_cuda=with_cuda,
        include_dirs = include_path,
        extra_link_args = [
            make_relative_rpath(lib_path),
            ENCODING_LIB,
        ],
Hang Zhang's avatar
init  
Hang Zhang committed
66
67
68
69
)

if __name__ == '__main__':
    ffi.build()