Split Training JNI targets to new 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.

Author: georgehong

extension/android/jni/BUCK

@@ -28,7 +28,7 @@ non_fbcode_target(_kind = executorch_generated_lib,
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_jni",
-    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
+    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS,
     soname = "libexecutorch.$(ext)",
@@ -39,19 +39,38 @@ non_fbcode_target(_kind = fb_android_cxx_library,
         "//fbandroid/libraries/fbjni:fbjni",
         "//fbandroid/native/fb:fb",
         "//third-party/glog:glog",
-        "//xplat/executorch/extension/data_loader:file_data_loader_static",
         "//xplat/executorch/extension/module:module_static",
         "//xplat/executorch/extension/runner_util:inputs_static",
         "//xplat/executorch/extension/tensor:tensor_static",
         "//xplat/executorch/extension/threadpool:threadpool_static",
-        "//xplat/executorch/extension/training/module:training_module_static",
-        "//xplat/executorch/extension/training/optimizer:sgd_static",
         third_party_dep("cpuinfo"),
     ],
 )
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_jni_full",
+    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
+    allow_jni_merging = False,
+    compiler_flags = ET_JNI_COMPILER_FLAGS,
+    soname = "libexecutorch.$(ext)",
+    visibility = ["PUBLIC"],
+    deps = [
+        ":jni_headers",
+        ":log_provider_static",
+        ":generated_op_lib_optimized_static",
+        "//fbandroid/libraries/fbjni:fbjni",
+        "//fbandroid/native/fb:fb",
+        "//third-party/glog:glog",
+        "//xplat/executorch/backends/xnnpack:xnnpack_backend_static",
+        "//xplat/executorch/extension/module:module_static",
+        "//xplat/executorch/extension/runner_util:inputs_static",
+        "//xplat/executorch/extension/tensor:tensor_static",
+        "//xplat/executorch/kernels/quantized:generated_lib_static",
+    ],
+)
+
+non_fbcode_target(_kind = fb_android_cxx_library,
+    name = "executorch_training_jni",
     srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS,
@@ -81,7 +100,6 @@ non_fbcode_target(_kind = fb_android_cxx_library,
         "jni_layer.cpp",
         "jni_layer_llama.cpp",
         "jni_layer_runtime.cpp",
-        "jni_layer_training.cpp",
     ],
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS + [
@@ -96,16 +114,13 @@ non_fbcode_target(_kind = fb_android_cxx_library,
         "//fbandroid/native/fb:fb",
         "//third-party/glog:glog",
         "//xplat/executorch/backends/xnnpack:xnnpack_backend_static",
-        "//xplat/executorch/extension/data_loader:file_data_loader_static",
         "//xplat/executorch/examples/models/llama/runner:runner_static",
         "//xplat/executorch/examples/models/llava/runner:runner_static",
         "//xplat/executorch/extension/module:module_static",
         "//xplat/executorch/extension/runner_util:inputs_static",
         "//xplat/executorch/extension/tensor:tensor_static",
         "//xplat/executorch/extension/threadpool:cpuinfo_utils_static",
         "//xplat/executorch/extension/threadpool:threadpool_static",
-        "//xplat/executorch/extension/training/module:training_module_static",
-        "//xplat/executorch/extension/training/optimizer:sgd_static",
     ],
 )
 
@@ -128,10 +143,6 @@ runtime.export_file(
     name = "jni_layer_runtime.cpp",
 )
 
-runtime.export_file(
-    name = "jni_layer_training.cpp",
-)
-
 runtime.cxx_library(
     name = "jni_headers",
     exported_headers = [

extension/android/jni/jni_layer.cpp

@@ -496,7 +496,7 @@ extern void register_natives_for_llm();
 void register_natives_for_llm() {}
 #endif
 extern void register_natives_for_runtime();
-extern void register_natives_for_training();
+ET_WEAK void register_natives_for_training() {};
 
 JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
   return facebook::jni::initialize(vm, [] {

extension/android/jni/selective_jni.buck.bzl

@@ -9,7 +9,6 @@ def selective_jni_target(name, deps, srcs = [], soname = "libexecutorch.$(ext)")
         name = name,
         srcs = [
             "//xplat/executorch/extension/android/jni:jni_layer.cpp",
-            "//xplat/executorch/extension/android/jni:jni_layer_training.cpp",
             "//xplat/executorch/extension/android/jni:jni_layer_runtime.cpp",
         ] + srcs,
         allow_jni_merging = False,
@@ -22,13 +21,10 @@ def selective_jni_target(name, deps, srcs = [], soname = "libexecutorch.$(ext)")
             "//third-party/glog:glog",
             "//xplat/executorch/extension/android/jni:jni_headers",
             "//xplat/executorch/extension/android/jni:log_provider_static",
-            "//xplat/executorch/extension/data_loader:file_data_loader_static",
             "//xplat/executorch/extension/module:module_static",
             "//xplat/executorch/extension/runner_util:inputs_static",
             "//xplat/executorch/extension/tensor:tensor_static",
             "//xplat/executorch/extension/threadpool:threadpool_static",
-            "//xplat/executorch/extension/training/module:training_module_static",
-            "//xplat/executorch/extension/training/optimizer:sgd_static",
             third_party_dep("cpuinfo"),
         ] + deps,
     )