TFLite selective builds using TF ops (flex delegate) for embedded linux on aarch64
Hi @JoshPPrieto, I wouldn't call it expected, we always appreciate it when the users share any research they have already done as it usually does help us. Feel free to make a PR that helps resolve the issue... if it affects other builds we'll have to analyze whether the effect is a regression or not but it'll be easier for us to examine when the PR is created.

@terryheo, can you please take a look? Thanks. Hi again, sorry for the delay.

As the changes contain some issues with latest master (remaining double registrations), rather than raising a PR (I don't think it complies with contribution guidelines), I'm sharing the commit (https://github.com/nxp-imx/tensorflow/commit/bd29abc2798e1bb28ac571bbb403d980eac52df0) to continue the discussion.

Please, @pkgoogle @terryheo, let me know if the PR is still preferred or something else is needed for reproduction/analysis.



### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.10 ... 2.15

### Custom code

No

### OS platform and distribution

Linux Ubuntu 22.04

### Mobile device

Yocto based Linux running kernel 6.1.x

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Hi,

From what it is described in [Select TensorFlow operators](https://www.tensorflow.org/lite/guide/ops_select) and [Reduce TensorFlow Lite binary size](https://www.tensorflow.org/lite/guide/reduce_binary_size), it is possible to generate reduced size binaries (minimal TFLite runtime + specific Flex ops) for Android, and it is also described how to build custom C/C++ shared libraries containing the Flex ops that are part of the given models during the build process.

When building the shared libs with models containing flex ops and `elinux_aarch64` as config, the expected artifacts are built with the full size, including all the TF ops rather than selecting the ones from the model. The behavior should be reproduced when trying to build benchmark-model with flex (or even the standalone shared lib):

tmp/BUILD (change `init_tensorflow` visibility if needed, `custom-model.tflite` should be a model containing at least 1 flex op)
```
load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_cc_binary", "clean_dep")
load("@org_tensorflow//tensorflow/lite:build_def.bzl", "tflite_copts", "tflite_copts_warnings", "tflite_linkopts")
load("@org_tensorflow//tensorflow/lite/delegates/flex:build_def.bzl", "tflite_flex_shared_library")

tflite_flex_shared_library(
    name = "tensorflowlite_flex_dynamic",
    models=[
       ":custom-model.tflite",
    ],
)

cc_import(
    name = "libtensorflowlite_flex_dynamic",
    shared_library = ":tensorflowlite_flex_dynamic",
)

tf_cc_binary(
    name = "benchmark_model_plus_flex_dynamic",
    srcs = [
        "//tensorflow/lite/tools/benchmark:benchmark_plus_flex_main.cc",
    ],
    copts = tflite_copts() + tflite_copts_warnings(),
    linkopts = tflite_linkopts(),
    deps = [
        ":libtensorflowlite_flex_dynamic",
        "//tensorflow/lite/tools/benchmark:benchmark_tflite_model_lib",
        "//tensorflow/lite/testing:init_tensorflow",
        "//tensorflow/lite/tools:logging",
    ],
)

```

Build command

```
bazel build -c opt \
	--cxxopt=`--std=c++17` \
	--config=monolithic \
	--config=elinux_aarch64 \
	--host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
	--verbose_failures  \
	//tmp:tensorflowlite_flex_dynamic
```


I'd like to know if building it with the full size is the expected behavior. If yes, I'd like to know if enabling it is part of the roadmap as it _**seems**_ like a low fruit hanging (it already works for same arch in Android). 

---

I experimented a little bit, partially succeeding (could build reduced size binaries by breaking plenty of bazel targets). I'm adding it to the issue with the hope that someone might find it helpful (disclaimer: take it with a pinch of salt as I've got skill issues with TF internals and the bazel build system).

+ Naively, find all targets that selectively depend on `android` or `mobile` configs and ensure that a branch includes `elinux_aarch64` (usually the targets are `portable_tensorflow_lib` and `portable_tensorflow_lib_lite`). This can be done adding `elinux_aarch64` as part of `mobile` configuration and/or `if_*` functions wrapping `select()`. Long story short, replicate dependencies used for `android` config (when makes sense). Ensure that `IS_MOBILE_PLATFORM` is defined for `portable_tensorflow_lib_lite`.

+ Caveats: the python tool to generate the `ops_to_register.h` header (`print_selective_registration_header`) is built natively, and there are some conflicts. One of them is a double registration issue (maybe related to some diamond dependency problem?). The other is related to some missing shared libs used by the tool. The only way I found to fix this is changing dependencies of targets from `tensorflow/python/*/BUILD`. Of course, this affects other dependent targets. I also tried to duplicate targets with modified names to avoid breaking the original ones, but I was unsuccessful.

If relevant, I can try to share the full patch. Just remember the changes will affect building other targets.

Some size results for `libtensorflowlite_flex.so`

(opt, all) -> 99 Mb
(opt, selected) -> 7 Mb
(dbg, all) -> 5.2 Gb
(dbg, selected) -> 2.1 Gb







### Standalone code to reproduce the issue

```shell
Check *current behavior* section. Generate the tmp/BUILD file, ensure a model with flex ops is present, and name of the `models` argument in `tflite_flex_shared_library` is correct.
```


### Relevant log output

_No response_