Skip to content

Commit 50f5032

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 956124e commit 50f5032

File tree

5 files changed

+40
-19
lines changed

5 files changed

+40
-19
lines changed

‎extension/android/CMakeLists.txt‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ set(executorch_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../lib/cmake/ExecuTorch)
6464
find_package(executorch CONFIG REQUIRED)
6565
target_link_options_shared_lib(executorch)
6666

67-
add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp jni/jni_layer_training.cpp)
67+
add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp)
6868

6969
set(link_libraries)
7070
list(
@@ -77,7 +77,6 @@ list(
7777
extension_runner_util
7878
extension_tensor
7979
extension_threadpool
80-
extension_training
8180
fbjni
8281
)
8382

@@ -147,6 +146,12 @@ if(EXECUTORCH_JNI_CUSTOM_LIBRARY)
147146
)
148147
endif()
149148

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

‎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_TRAINING_JNI",
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_TRAINING_JNI
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
)

‎scripts/build_android_library.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ build_android_native_library() {
8484
-DNEURON_BUFFER_ALLOCATOR_LIB="$NEURON_BUFFER_ALLOCATOR_LIB" \
8585
-DEXECUTORCH_BUILD_KERNELS_CUSTOM="${EXECUTORCH_BUILD_EXTENSION_LLM:-ON}" \
8686
-DEXECUTORCH_BUILD_LLAMA_JNI="${EXECUTORCH_BUILD_EXTENSION_LLM:-ON}" \
87+
-DEXECUTORCH_BUILD_TRAINING_JNI="${EXECUTORCH_BUILD_EXTENSION_LLM:-ON}" \
8788
-DSUPPORT_REGEX_LOOKAHEAD=ON \
8889
-DCMAKE_BUILD_TYPE="${EXECUTORCH_CMAKE_BUILD_TYPE}" \
8990
-B"${CMAKE_OUT}"/extension/android

0 commit comments

Comments
 (0)