Include native C symbols and headers for TFLite AARs

RELNOTES=TensorFlow Lite Android AARs now include the C headers
and APIs required to use TFLite from native code.
PiperOrigin-RevId: 284294408
Change-Id: I10052adba8355cb0d980728aca518efad1187272

tensorflow/lite/delegates/gpu/BUILD

@@ -6,7 +6,10 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
-exports_files(["metal_delegate.h"])
+exports_files([
+    "gpu_delegate.h",
+    "metal_delegate.h",
+])
 
 # Primary purpose of this config is to replace ::util::Status with our custom
 # light implementation ::tflite::gpu::StatusLite to reduce binary size.  Besides

tensorflow/lite/java/BUILD

@@ -25,6 +25,12 @@ JAVA_SRCS = glob([
 aar_with_jni(
     name = "tensorflow-lite",
     android_library = ":tensorflowlite",
+    headers = [
+        "//tensorflow/lite:builtin_ops.h",
+        "//tensorflow/lite/c:c_api.h",
+        "//tensorflow/lite/c:c_api_experimental.h",
+        "//tensorflow/lite/c:common.h",
+    ],
 )
 
 # EXPERIMENTAL: AAR target for using TensorFlow ops with TFLite. Note that this
@@ -49,6 +55,9 @@ aar_with_jni(
 aar_with_jni(
     name = "tensorflow-lite-gpu",
     android_library = ":tensorflowlite_gpu",
+    headers = [
+        "//tensorflow/lite/delegates/gpu:delegate.h",
+    ],
 )
 
 android_library(
@@ -349,7 +358,12 @@ cc_library(
 
 tflite_jni_binary(
     name = "libtensorflowlite_jni.so",
+    linkscript = ":tflite_version_script.lds",
     deps = [
+        # Note that we explicitly include the C API here for convenience, as it
+        # allows bundling of the C lib w/ AAR distribution.
+        "//tensorflow/lite/c:c_api",
+        "//tensorflow/lite/c:c_api_experimental",
         "//tensorflow/lite/delegates/nnapi/java/src/main/native",
         "//tensorflow/lite/java/src/main/native",
     ],
@@ -366,6 +380,7 @@ tflite_jni_binary(
 # EXPERIMENTAL: Native target that supports GPU acceleration.
 tflite_jni_binary(
     name = "libtensorflowlite_gpu_jni.so",
+    linkscript = ":gpu_version_script.lds",
     deps = [
         "//tensorflow/lite/delegates/gpu/java/src/main/native",
     ],

tensorflow/lite/java/aar_with_jni.bzl

@@ -5,7 +5,8 @@ load("@build_bazel_rules_android//android:rules.bzl", "android_binary")
 def aar_with_jni(
         name,
         android_library,
-        headers = None):
+        headers = None,
+        flatten_headers = False):
     """Generates an Android AAR given an Android library target.
 
     Args:
@@ -16,6 +17,7 @@ def aar_with_jni(
       headers: Optional list of headers that will be included in the
           generated .aar file. This is useful for distributing self-contained
           .aars with native libs that can be used directly by native clients.
+      flatten_headers: Whether to flatten the output paths of included headers.
     """
 
     # Generate dummy AndroidManifest.xml for dummy apk usage
@@ -68,9 +70,15 @@ zip -r $$origdir/$(location :{1}.aar) jni/*/*.so
         mkdir headers
         """
         for src in headers:
+            if flatten_headers:
                 cmd += """
                     cp -rL $$origdir/$(location {0}) headers/$$(basename $(location {0}))
                 """.format(src)
+            else:
+                cmd += """
+                    mkdir -p headers/$$(dirname $(location {0}))
+                    cp -rL $$origdir/$(location {0}) headers/$(location {0})
+                """.format(src)
         cmd += "zip -r $$origdir/$(location :{0}.aar) headers".format(name)
 
     native.genrule(

tensorflow/lite/java/gpu_version_script.lds

@@ -0,0 +1,12 @@
+VERS_1.0 {
+  # Export JNI and native C symbols.
+  global:
+    Java_*;
+    JNI_OnLoad;
+    JNI_OnUnload;
+    TfLiteGpu*;
+
+  # Hide everything else.
+  local:
+    *;
+};

tensorflow/lite/java/tflite_version_script.lds

@@ -0,0 +1,12 @@
+VERS_1.0 {
+  # Export JNI and native C symbols.
+  global:
+    Java_*;
+    JNI_OnLoad;
+    JNI_OnUnload;
+    TfLite*;
+
+  # Hide everything else.
+  local:
+    *;
+};