Skip to content

Commit 015954e

Browse files
larryliu0820facebook-github-bot
authored andcommitted
Lookup for libtorch.so and create buck2 target
Summary: `libtorch.so` is needed by custom ops as well as ATen mode. This diff adds a script to find the path of installed PyTorch directory in Python library and then look for `libtorch.so`. This script also copy header files into buck out so that they can be included by buck targets depending on the genrule. Note that this only works if user install pytorch properly. For example: ## libtorch.so My environment has `libtorch.so` installed in: ``` /opt/conda/envs/executorch/lib/python3.10/site-packages/torch/lib/libtorch.so ``` The script `link_torch.sh` creates a symbolic link in buck out: ``` lrwxrwxrwx 1 vscode vscode 77 Aug 10 07:43 /workspaces/executorch/buck-out/v2/gen/root/524f8da68ea2a374/third-party/__libtorch_gen__/out/libtorch.so -> /opt/conda/envs/executorch/lib/python3.10/site-packages/torch/lib/libtorch.so ``` Then we wrap it using a `prebuilt_cxx_library` macro. ## headers The script then go ahead and copy the headers under: ``` /opt/conda/envs/executorch/lib/python3.10/site-packages/torch/include ``` to buck-out: ``` /workspaces/executorch/buck-out/v2/gen/root/524f8da68ea2a374/third-party/__torch_headers_gen__/out ``` In the library definition we also need to add `-I<header_dir>` external preprocessor flags, so when we link `libtorch` we are also including those header directories. As a result, we will be able to link `//third-party:libtorch` to any C++ source file. We then are able to build `custom_ops_generated_lib` which registers custom ops using PyTorch C++ API. ## linker flag We need to add `-Wl,-rpath,$(location :libtorch_dir)` external linker flag, for any entity doing a `dlopen` to be able to find this shared library. Reviewed By: kimishpatel Differential Revision: D48221853 fbshipit-source-id: 0c7706c6f33f36d21212f965bd3c23aa9b2da0af
1 parent 3b0533f commit 015954e

File tree

7 files changed

+73
-7
lines changed

7 files changed

+73
-7
lines changed

examples/custom_ops/custom_ops.yaml

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,9 @@
77
- func: my_ops::mul3.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)
88
kernels:
99
- arg_meta: null
10-
kernel_name: custom::mul3_out_impl # sub-namespace native:: is auto-added
10+
kernel_name: custom::mul3_out_impl # custom_op_1.cpp, sub-namespace native:: is auto-added
11+
12+
- func: my_ops::mul4.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)
13+
kernels:
14+
- arg_meta: null
15+
kernel_name: custom::mul4_out_impl # custom_op_2.cpp

runtime/core/exec_aten/targets.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,5 +20,5 @@ def define_common_targets():
2020
"@EXECUTORCH_CLIENTS",
2121
],
2222
exported_deps = ["//executorch/runtime/core:tensor_shape_dynamism"] + ([] if aten_mode else ["//executorch/runtime/core/portable_type:portable_type"]),
23-
exported_external_deps = ["torch-core-cpp"] if aten_mode else [],
23+
exported_external_deps = ["libtorch"] if aten_mode else [],
2424
)

runtime/core/exec_aten/util/scalar_type_util.h

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -318,9 +318,10 @@ inline bool isComplexType(exec_aten::ScalarType t) {
318318
}
319319

320320
inline bool isBitsType(exec_aten::ScalarType t) {
321-
return t == ScalarType::Bits1x8 || t == ScalarType::Bits2x4 ||
322-
t == ScalarType::Bits4x2 || t == ScalarType::Bits8 ||
323-
t == ScalarType::Bits16;
321+
return t == exec_aten::ScalarType::Bits1x8 ||
322+
t == exec_aten::ScalarType::Bits2x4 ||
323+
t == exec_aten::ScalarType::Bits4x2 ||
324+
t == exec_aten::ScalarType::Bits8 || t == exec_aten::ScalarType::Bits16;
324325
}
325326

326327
inline bool isQIntType(exec_aten::ScalarType t) {

runtime/core/exec_aten/util/targets.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ def define_common_targets():
2424
exported_deps = [
2525
"//executorch/runtime/core:core",
2626
] + [] if aten_mode else ["//executorch/runtime/core/portable_type:scalar_type"],
27-
exported_external_deps = ["torch-core-cpp"] if aten_mode else [],
27+
exported_external_deps = ["libtorch"] if aten_mode else [],
2828
)
2929

3030
runtime.cxx_library(

shim/xplat/executorch/build/env_interface.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ _EXTERNAL_DEPS = {
3030
# Commandline flags library
3131
"gflags": "//third-party:gflags",
3232
"gmock": [], # TODO(larryliu0820): Add support
33-
"libtorch": [], # TODO(larryliu0820): Add support
33+
"libtorch": "//third-party:libtorch",
3434
"prettytable": "//third-party:prettytable",
3535
"pybind11": [], # TODO(larryliu0820): Add support
3636
# Core C++ PyTorch functionality like Tensor and ScalarType.

third-party/TARGETS

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ load("@fbsource//xplat/executorch/build/runtime_wrapper.bzl", "runtime")
22
load(":gflags.bzl", "define_gflags")
33
load(":glob_defs.bzl", "subdir_glob")
44
load(":prebuilt_python_defs.bzl", "add_prebuilt_python_library_targets")
5+
load("@prelude//rules.bzl", "prebuilt_cxx_library")
56

67
define_gflags()
78

@@ -204,3 +205,38 @@ runtime.cxx_binary(
204205
link_style = "static",
205206
_is_external_target = True,
206207
)
208+
209+
runtime.genrule(
210+
name = "libtorch_gen",
211+
out = "libtorch.so",
212+
srcs = ["link_torch.sh"],
213+
bash = 'bash $SRCS torch/lib/libtorch.so ${OUT}',
214+
)
215+
216+
runtime.genrule(
217+
name = "torch_headers_gen",
218+
out = "include",
219+
srcs = ["link_torch.sh"],
220+
bash = 'bash $SRCS torch/include ${OUT}',
221+
)
222+
223+
runtime.genrule(
224+
name = "libtorch_dir",
225+
out = "lib",
226+
srcs = ["link_torch.sh"],
227+
bash = 'bash $SRCS torch/lib ${OUT}',
228+
)
229+
230+
prebuilt_cxx_library(
231+
name = "libtorch",
232+
shared_lib = ":libtorch_gen",
233+
soname = "libtorch.so",
234+
exported_preprocessor_flags = [
235+
"-D_GLIBCXX_USE_CXX11_ABI=0", # `libtorch` is built without CXX11_ABI so any target depends on it need to use the same build config.
236+
"-I$(location :torch_headers_gen)", # include header directories
237+
"-I$(location :torch_headers_gen)/torch/csrc/api/include", # include header directories
238+
],
239+
exported_linker_flags = ["-Wl,-rpath,$(location :libtorch_dir)"], # define rpath to locate shared library
240+
exported_headers = [":torch_headers_gen"],
241+
visibility = ["PUBLIC"],
242+
)

third-party/link_torch.sh

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
# Links the path of `libtorch.so` or include library installed in python
9+
# library, to the output of buck, if pytorch is properly installed.
10+
# This util can be used by buck2 build.
11+
12+
set -e
13+
14+
LIB=$(python3 -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
15+
16+
SUBPATH=$1
17+
OUT=$2
18+
if [[ -f "$LIB/$SUBPATH" ]] || [[ -d "$LIB/$SUBPATH" ]];
19+
then
20+
ln -s "$LIB/$SUBPATH" "$OUT"
21+
else
22+
echo "Error: $LIB/$SUBPATH doesn't exist"
23+
exit 1
24+
fi

0 commit comments

Comments
 (0)