Always use two XNNPACK Partitioners (#5573)

Summary:
Pull Request resolved: #5573

This changes the default behavior. Helps prefill ~20%, hurts decode ~7%.

As a next step, I will try to debug more into perf regression on decode and if anything more we can get on prefill by tuning xnnpack thread dispatcher for gemm, gemv, mul, add, sigmoid, and sub.

**On my local (unreliable) S23** -

* Vanilla:

```
dm1q:/data/local/tmp/llama $ ./llama_main_release  \
--model_path ./llama_gs32_vanilla.pte  \
--tokenizer_path ./tokenizer.bin \
--seq_len=128 \
--prompt="${prompt}"

[...]
I 00:00:22.188618 executorch:stats.h:84]        Prompt Tokens: 44    Generated Tokens: 83
I 00:00:22.188621 executorch:stats.h:90]        Model Load Time:                12.922000 (seconds)
I 00:00:22.188624 executorch:stats.h:100]       Total inference time:           9.252000 (seconds)               Rate:  8.971033 (tokens/second)
I 00:00:22.188627 executorch:stats.h:108]               Prompt evaluation:      1.740000 (seconds)               Rate:  25.287356 (tokens/second)
I 00:00:22.188630 executorch:stats.h:119]               Generated 83 tokens:    7.512000 (seconds)               Rate:  11.048988 (tokens/second)
I 00:00:22.188632 executorch:stats.h:127]       Time to first generated token:  1.740000 (seconds)
I 00:00:22.188634 executorch:stats.h:134]       Sampling time over 127 tokens:  0.015000 (seconds)
[...]
```

* Two partition (2part)
```
dm1q:/data/local/tmp/llama $ ./llama_main_release \
 --model_path ./llama_gs32_2part.pte  \  # New PTE
--tokenizer_path ./tokenizer.bin \
--seq_len=128 \
--prompt="${prompt}"

[...]
I 00:00:22.205058 executorch:stats.h:84]        Prompt Tokens: 44    Generated Tokens: 83
I 00:00:22.205061 executorch:stats.h:90]        Model Load Time:                12.876000 (seconds)
I 00:00:22.205063 executorch:stats.h:100]       Total inference time:           9.323000 (seconds)               Rate:  8.902714 (tokens/second)
I 00:00:22.205067 executorch:stats.h:108]               Prompt evaluation:      1.549000 (seconds)               Rate:  28.405423 (tokens/second)
I 00:00:22.205070 executorch:stats.h:119]               Generated 83 tokens:    7.774000 (seconds)               Rate:  10.676614 (tokens/second)
I 00:00:22.205073 executorch:stats.h:127]       Time to first generated token:  1.549000 (seconds)
I 00:00:22.205075 executorch:stats.h:134]       Sampling time over 127 tokens:  0.029000 (seconds)
[...]
```

**Similar results on AiBench OnePlus12**,
* Vanilla, AiBench Links: [gs=32](https://www.internalfb.com/intern/aibench/details/114258284562772), [gs=256](https://www.internalfb.com/intern/aibench/details/438103192423336)
```
# gs=32

I 00:00:21.792659 executorch:stats.h:84] 	Prompt Tokens: 5    Generated Tokens: 118
I 00:00:21.792721 executorch:stats.h:90] 	Model Load Time:		11.666000 (seconds)
I 00:00:21.792754 executorch:stats.h:100] 	Total inference time:		10.109000 (seconds)		 Rate: 	11.672767 (tokens/second)
I 00:00:21.792778 executorch:stats.h:108] 		Prompt evaluation:	0.365000 (seconds)		 Rate: 	13.698630 (tokens/second)
I 00:00:21.792799 executorch:stats.h:119] 		Generated 118 tokens:	9.744000 (seconds)		 Rate: 	12.110016 (tokens/second)
I 00:00:21.792818 executorch:stats.h:127] 	Time to first generated token:	0.365000 (seconds)
I 00:00:21.792837 executorch:stats.h:134] 	Sampling time over 123 tokens:	0.008000 (seconds)
```

* Two partition,  AiBench Links:  [gs=32](https://www.internalfb.com/intern/aibench/details/852029802754424), [gs=256](https://www.internalfb.com/intern/aibench/details/491722732991273)
```
# gs=32

I 00:00:22.584271 executorch:stats.h:84] 	Prompt Tokens: 5    Generated Tokens: 118
I 00:00:22.584336 executorch:stats.h:90] 	Model Load Time:		11.610000 (seconds)
I 00:00:22.584367 executorch:stats.h:100] 	Total inference time:		10.960000 (seconds)		 Rate: 	10.766423 (tokens/second)
I 00:00:22.584389 executorch:stats.h:108] 		Prompt evaluation:	0.286000 (seconds)		 Rate: 	17.482517 (tokens/second)
I 00:00:22.584409 executorch:stats.h:119] 		Generated 118 tokens:	10.674000 (seconds)		 Rate: 	11.054900 (tokens/second)
I 00:00:22.584428 executorch:stats.h:127] 	Time to first generated token:	0.286000 (seconds)
I 00:00:22.584446 executorch:stats.h:134] 	Sampling time over 123 tokens:	0.013000 (seconds)
```

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: mcr229

Differential Revision: D63271101

fbshipit-source-id: 1680c8d97ec5ac15011eae8a6d75a003bd0354b4

Author: digantdesai

.ci/scripts/test_llama.sh

@@ -188,7 +188,7 @@ EXPORTED_MODEL_NAME="${EXPORTED_MODEL_NAME}.pte"
 echo "Exporting ${EXPORTED_MODEL_NAME}"
 EXPORT_ARGS="-c ${CHECKPOINT_FILE_NAME} -p ${PARAMS} -d ${DTYPE} -n ${EXPORTED_MODEL_NAME} -kv"
 if [[ "${XNNPACK}" == "ON" ]]; then
-  EXPORT_ARGS="${EXPORT_ARGS} -X -qmode 8da4w -G 128"
+  EXPORT_ARGS="${EXPORT_ARGS} -X --xnnpack-extended-ops -qmode 8da4w -G 128"
 fi
 if [[ "${CUSTOM}" == "ON" ]]; then
   EXPORT_ARGS="${EXPORT_ARGS} --use_sdpa_with_kv_cache"

examples/models/llama2/export_llama_lib.py

@@ -297,7 +297,17 @@ def build_args_parser() -> argparse.ArgumentParser:
 
     parser.add_argument("-2", "--fairseq2", action="store_true")
     parser.add_argument("-v", "--verbose", action="store_true")
-    parser.add_argument("-X", "--xnnpack", action="store_true")
+    parser.add_argument(
+        "-X",
+        "--xnnpack",
+        action="store_true",
+        help="Delegate to DQLinear ops to the xnnpack backend",
+    )
+    parser.add_argument(
+        "--xnnpack-extended-ops",
+        action="store_true",
+        help="Delegate more operators beyond DQLinear to the xnnpack backend. Requires -X or --xnnpack to be set.",
+    )
     parser.add_argument("-V", "--vulkan", action="store_true")
     parser.add_argument("--mps", action="store_true")
     parser.add_argument("--coreml", action="store_true")
@@ -546,12 +556,24 @@ def _export_llama(modelname, args) -> LLMEdgeManager:  # noqa: C901
 
     # to_backend
     partitioners = []
-    if pt2e_quant_params is not None and pt2e_quant_params.quantize_linear is not None:
-        partitioners.append(get_xnnpack_partitioner())
+
+    # Order matters here, dynamic quantization should be applied first when both xnnpack and xnnpack_extended_ops are enabled
+    if (
+        pt2e_quant_params is not None and pt2e_quant_params.quantize_linear is not None
+    ) or (args.xnnpack):
+        partitioners.append(
+            get_xnnpack_partitioner(dynamic_quant_only_partitioner=True)
+        )
+
+        # force xnnpack to be true if pt2e_quant_params is not None and args.xnnpack is False
+        args.xnnpack = True
         modelname = f"xnnpack_dq_{modelname}"
 
-    if args.xnnpack:
-        partitioners.append(get_xnnpack_partitioner())
+    if args.xnnpack_extended_ops:
+        assert args.xnnpack, "xnnpack_extended_ops requires xnnpack to be enabled"
+        partitioners.append(
+            get_xnnpack_partitioner(dynamic_quant_only_partitioner=False)
+        )
         modelname = f"xnnpack_{modelname}"
 
     if args.vulkan:
@@ -598,6 +620,10 @@ def _export_llama(modelname, args) -> LLMEdgeManager:  # noqa: C901
                 shares=args.num_sharding,
             )
 
+    logging.info("Lowering model using following partitioner(s): ")
+    for partitioner in partitioners:
+        logging.info(f"--> {partitioner.__class__.__name__}")
+
     if args.generate_etrecord:
         if not builder_exported_to_edge.edge_manager:
             raise ValueError("Unable to generate etrecord due to missing edge manager.")

extension/llm/export/partitioner_lib.py

@@ -7,16 +7,28 @@
 from typing import Optional
 
 
-def get_xnnpack_partitioner():
+def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True):
+    """
+    Returns the XNNPACK partitioner.
+
+    @arg dynamic_quant_only_partitioner:
+        This is enabled by default to keep BC.
+        If dynamic_quant_only_partitioner is True, then only dynamically quantized
+        linear layers will be partitioned.
+        Else, anything which can be will be partitioned greedily.
+    """
     from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
         XnnpackDynamicallyQuantizedPartitioner,
+        XnnpackPartitioner,
     )
 
-    # Following changes due to.
-    # 1. We need dynamically quantized partitioner for both pt2e_quantize options
-    #    as well as "qmode 8da4w" which is also dynamic quantizes linear layers.
-    # 2. XNNPACK partitioner seems to result in seg fault for non dqlinear ops.
-    return XnnpackDynamicallyQuantizedPartitioner()
+    if dynamic_quant_only_partitioner:
+        # Following changes due to.
+        # 1. We need dynamically quantized partitioner for both pt2e_quantize options
+        #    as well as "qmode 8da4w" which is also dynamic quantizes linear layers.
+        # 2. XNNPACK partitioner seems to result in seg fault for non dqlinear ops.
+        return XnnpackDynamicallyQuantizedPartitioner()
+    return XnnpackPartitioner()
 
 
 def get_vulkan_partitioner(