From ab6d91c5c2d9ffa2dd01695e48fdc2981f16085e Mon Sep 17 00:00:00 2001
From: Hansong <107070759+kirklandsign@users.noreply.github.com>
Date: Tue, 10 Sep 2024 14:50:15 -0700
Subject: [PATCH] Fix internal executorch_llama_jni

Differential Revision: D62458604

Pull Request resolved: https://github.com/pytorch/executorch/pull/5231
---
 extension/android/jni/BUCK | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/extension/android/jni/BUCK b/extension/android/jni/BUCK
index f7e7932a21..3c8f00b2bd 100644
--- a/extension/android/jni/BUCK
+++ b/extension/android/jni/BUCK
@@ -70,11 +70,16 @@ fb_android_cxx_library(
 
 fb_android_cxx_library(
     name = "executorch_llama_jni",
-    srcs = ["jni_layer_llama.cpp"],
+    srcs = [
+        "jni_layer.cpp",
+        "jni_layer_llama.cpp",
+    ],
+    headers = ["jni_layer_constants.h"],
     allow_jni_merging = False,
     compiler_flags = [
         "-frtti",
         "-fexceptions",
+        "-DEXECUTORCH_BUILD_LLAMA_JNI",
         "-Wno-format",
     ],
     soname = "libexecutorch.$(ext)",
@@ -83,8 +88,12 @@ fb_android_cxx_library(
         "//fbandroid/libraries/fbjni:fbjni",
         "//fbandroid/native/fb:fb",
         "//third-party/glog:glog",
+        "//xplat/executorch/backends/xnnpack:xnnpack_backend_static",
         "//xplat/executorch/examples/models/llama2/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",
     ],