Skip to content

Commit b4ca137

Browse files
georgehonggeorgehong
committed
Guard training JNI targets with training JNI flag and split BUCK target
Address comment on JNI binary size sensitivity. Rather than adding to the existing JNI buck targets, initially introduce a new executorch_training_jni target. Using EXECUTORCH_BUILD_TRAINING_JNI to further modularize JNI build
1 parent 04e3e76 commit b4ca137

File tree

5 files changed

+40
-20
lines changed

5 files changed

+40
-20
lines changed

‎extension/android/CMakeLists.txt‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ set_target_properties(fbjni PROPERTIES
6363
set(executorch_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../lib/cmake/ExecuTorch)
6464
target_link_options_shared_lib(executorch)
6565

66-
add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp jni/jni_layer_training.cpp)
66+
add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp)
6767

6868
set(link_libraries)
6969
list(
@@ -76,7 +76,6 @@ list(
7676
extension_runner_util
7777
extension_tensor
7878
extension_threadpool
79-
extension_training
8079
fbjni
8180
)
8281

@@ -146,6 +145,12 @@ if(EXECUTORCH_JNI_CUSTOM_LIBRARY)
146145
)
147146
endif()
148147

148+
if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
149+
target_sources(executorch_jni PRIVATE jni/jni_layer_training.cpp jni/log.cpp)
150+
list(APPEND link_libraries extension_training)
151+
target_compile_definitions(executorch_jni PUBLIC EXECUTORCH_BUILD_EXTENSION_TRAINING=1)
152+
endif()
153+
149154
if(EXECUTORCH_BUILD_LLAMA_JNI)
150155
target_sources(executorch_jni PRIVATE jni/jni_layer_llama.cpp jni/log.cpp)
151156
list(APPEND link_libraries llama_runner llava_runner)

‎extension/android/executorch_android/android_test_setup.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ prepare_add() {
2121
}
2222

2323
prepare_xor() {
24-
pushd "${BASEDIR}/../../../extension/training/"
24+
pushd "${BASEDIR}/../../training/"
2525
python3 -m examples.XOR.export_model --outdir "${BASEDIR}/src/androidTest/resources/"
2626
mv "${BASEDIR}/src/androidTest/resources/xor.pte" "${BASEDIR}/src/androidTest/resources/xor_full.pte"
2727
python3 -m examples.XOR.export_model --outdir "${BASEDIR}/src/androidTest/resources/" --external

‎extension/android/jni/BUCK‎

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ non_fbcode_target(_kind = executorch_generated_lib,
2828

2929
non_fbcode_target(_kind = fb_android_cxx_library,
3030
name = "executorch_jni",
31-
srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
31+
srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
3232
allow_jni_merging = False,
3333
compiler_flags = ET_JNI_COMPILER_FLAGS,
3434
soname = "libexecutorch.$(ext)",
@@ -39,24 +39,45 @@ non_fbcode_target(_kind = fb_android_cxx_library,
3939
"//fbandroid/libraries/fbjni:fbjni",
4040
"//fbandroid/native/fb:fb",
4141
"//third-party/glog:glog",
42-
"//xplat/executorch/extension/data_loader:file_data_loader_static",
4342
"//xplat/executorch/extension/module:module_static",
4443
"//xplat/executorch/extension/runner_util:inputs_static",
4544
"//xplat/executorch/extension/tensor:tensor_static",
4645
"//xplat/executorch/extension/threadpool:threadpool_static",
47-
"//xplat/executorch/extension/training/module:training_module_static",
48-
"//xplat/executorch/extension/training/optimizer:sgd_static",
4946
third_party_dep("cpuinfo"),
5047
],
5148
)
5249

5350
non_fbcode_target(_kind = fb_android_cxx_library,
5451
name = "executorch_jni_full",
55-
srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
52+
srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
5653
allow_jni_merging = False,
5754
compiler_flags = ET_JNI_COMPILER_FLAGS,
5855
soname = "libexecutorch.$(ext)",
5956
visibility = ["PUBLIC"],
57+
deps = [
58+
":jni_headers",
59+
":log_provider_static",
60+
":generated_op_lib_optimized_static",
61+
"//fbandroid/libraries/fbjni:fbjni",
62+
"//fbandroid/native/fb:fb",
63+
"//third-party/glog:glog",
64+
"//xplat/executorch/backends/xnnpack:xnnpack_backend_static",
65+
"//xplat/executorch/extension/module:module_static",
66+
"//xplat/executorch/extension/runner_util:inputs_static",
67+
"//xplat/executorch/extension/tensor:tensor_static",
68+
"//xplat/executorch/kernels/quantized:generated_lib_static",
69+
],
70+
)
71+
72+
non_fbcode_target(_kind = fb_android_cxx_library,
73+
name = "executorch_training_jni",
74+
srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
75+
allow_jni_merging = False,
76+
compiler_flags = ET_JNI_COMPILER_FLAGS + [
77+
"-DEXECUTORCH_BUILD_EXTENSION_TRAINING",
78+
],
79+
soname = "libexecutorch.$(ext)",
80+
visibility = ["PUBLIC"],
6081
deps = [
6182
":jni_headers",
6283
":log_provider_static",
@@ -81,7 +102,6 @@ non_fbcode_target(_kind = fb_android_cxx_library,
81102
"jni_layer.cpp",
82103
"jni_layer_llama.cpp",
83104
"jni_layer_runtime.cpp",
84-
"jni_layer_training.cpp",
85105
],
86106
allow_jni_merging = False,
87107
compiler_flags = ET_JNI_COMPILER_FLAGS + [
@@ -96,16 +116,13 @@ non_fbcode_target(_kind = fb_android_cxx_library,
96116
"//fbandroid/native/fb:fb",
97117
"//third-party/glog:glog",
98118
"//xplat/executorch/backends/xnnpack:xnnpack_backend_static",
99-
"//xplat/executorch/extension/data_loader:file_data_loader_static",
100119
"//xplat/executorch/examples/models/llama/runner:runner_static",
101120
"//xplat/executorch/examples/models/llava/runner:runner_static",
102121
"//xplat/executorch/extension/module:module_static",
103122
"//xplat/executorch/extension/runner_util:inputs_static",
104123
"//xplat/executorch/extension/tensor:tensor_static",
105124
"//xplat/executorch/extension/threadpool:cpuinfo_utils_static",
106125
"//xplat/executorch/extension/threadpool:threadpool_static",
107-
"//xplat/executorch/extension/training/module:training_module_static",
108-
"//xplat/executorch/extension/training/optimizer:sgd_static",
109126
],
110127
)
111128

@@ -128,10 +145,6 @@ runtime.export_file(
128145
name = "jni_layer_runtime.cpp",
129146
)
130147

131-
runtime.export_file(
132-
name = "jni_layer_training.cpp",
133-
)
134-
135148
runtime.cxx_library(
136149
name = "jni_headers",
137150
exported_headers = [

‎extension/android/jni/jni_layer.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,7 +496,13 @@ extern void register_natives_for_llm();
496496
void register_natives_for_llm() {}
497497
#endif
498498
extern void register_natives_for_runtime();
499+
500+
#ifdef EXECUTORCH_BUILD_EXTENSION_TRAINING
499501
extern void register_natives_for_training();
502+
#else
503+
// No op if we don't build training JNI
504+
void register_natives_for_training() {}
505+
#endif
500506

501507
JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
502508
return facebook::jni::initialize(vm, [] {

‎extension/android/jni/selective_jni.buck.bzl‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ def selective_jni_target(name, deps, srcs = [], soname = "libexecutorch.$(ext)")
99
name = name,
1010
srcs = [
1111
"//xplat/executorch/extension/android/jni:jni_layer.cpp",
12-
"//xplat/executorch/extension/android/jni:jni_layer_training.cpp",
1312
"//xplat/executorch/extension/android/jni:jni_layer_runtime.cpp",
1413
] + srcs,
1514
allow_jni_merging = False,
@@ -22,13 +21,10 @@ def selective_jni_target(name, deps, srcs = [], soname = "libexecutorch.$(ext)")
2221
"//third-party/glog:glog",
2322
"//xplat/executorch/extension/android/jni:jni_headers",
2423
"//xplat/executorch/extension/android/jni:log_provider_static",
25-
"//xplat/executorch/extension/data_loader:file_data_loader_static",
2624
"//xplat/executorch/extension/module:module_static",
2725
"//xplat/executorch/extension/runner_util:inputs_static",
2826
"//xplat/executorch/extension/tensor:tensor_static",
2927
"//xplat/executorch/extension/threadpool:threadpool_static",
30-
"//xplat/executorch/extension/training/module:training_module_static",
31-
"//xplat/executorch/extension/training/optimizer:sgd_static",
3228
third_party_dep("cpuinfo"),
3329
] + deps,
3430
)

0 commit comments

Comments
 (0)