build.py 1.68 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
backend  
Hang Zhang committed
17
18
19
20
21
22
23
24
25
26
# build kernel library
build_all_cmd = ['bash', 'encoding/make.sh']
if subprocess.call(build_all_cmd) != 0:
	sys.exit(1)

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

Hang Zhang's avatar
init  
Hang Zhang committed
27
28
29
30
package_base = os.path.dirname(torch.__file__)
this_file = os.path.dirname(os.path.realpath(__file__))

include_path = [os.path.join(os.environ['HOME'],'pytorch/torch/lib/THC'), 
Hang Zhang's avatar
Hang Zhang committed
31
32
								os.path.join(package_base,'lib/include/ENCODING'), 
								os.path.join(this_file,'encoding/src/')]
Hang Zhang's avatar
init  
Hang Zhang committed
33

Hang Zhang's avatar
Hang Zhang committed
34
35
36
37
if platform.system() == 'Darwin':
	ENCODING_LIB = os.path.join(package_base, 'lib/libENCODING.dylib')
else:
	ENCODING_LIB = os.path.join(package_base, 'lib/libENCODING.so')
Hang Zhang's avatar
init  
Hang Zhang committed
38

Hang Zhang's avatar
Hang Zhang committed
39
def make_relative_rpath(path):
Hang Zhang's avatar
tested  
Hang Zhang committed
40
41
42
43
44
	if platform.system() == 'Darwin':
		return '-Wl,-rpath,' + path
	else:
		return '-Wl,-rpath,' + path

Hang Zhang's avatar
init  
Hang Zhang committed
45
ffi = create_extension(
Hang Zhang's avatar
tested  
Hang Zhang committed
46
47
48
49
50
51
52
	'encoding._ext.encoding_lib',
	package=True,
	headers=headers,
	sources=sources,
	define_macros=defines,
	relative_to=__file__,
	with_cuda=with_cuda,
Hang Zhang's avatar
init  
Hang Zhang committed
53
		include_dirs = include_path,
Hang Zhang's avatar
Hang Zhang committed
54
55
56
57
		extra_link_args = [
			make_relative_rpath(os.path.join(package_base, 'lib')),
			ENCODING_LIB,
		],
Hang Zhang's avatar
init  
Hang Zhang committed
58
59
60
61
)

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