Merge branch 'main' into delete-exir-capture

Author: krisnaparahita

examples/models/llama/export_llama_lib.py

@@ -462,8 +462,11 @@ def build_args_parser() -> argparse.ArgumentParser:
     )
     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.",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Delegate more operators beyond DQLinear to the xnnpack backend. On by "
+        "default; pass --no-xnnpack-extended-ops to delegate DQLinear only. Requires "
+        "-X or --xnnpack to be set.",
     )
     parser.add_argument(
         "--xnnpack-enable-bf16",
@@ -1165,7 +1168,7 @@ def _to_edge_and_lower_llama_xnnpack(
     pt2e_quant_params,
     quantizers,
     quant_dtype,
-    xnnpack_extended_ops: bool = False,
+    xnnpack_extended_ops: bool = True,
     generate_etrecord: bool = False,
     verbose: bool = False,
     gen_tag_fn: Optional[Callable[[torch.fx.Node], Optional[str]]] = None,
@@ -1180,15 +1183,16 @@ def _to_edge_and_lower_llama_xnnpack(
         )
     )
 
-    modelname = f"xnnpack_dq_{modelname}"
-
     if xnnpack_extended_ops:
         partitioners.append(
             get_xnnpack_partitioner(
                 dynamic_quant_only_partitioner=False, enable_bf16=enable_bf16
             )
         )
-        modelname = f"xnnpack_{modelname}"
+
+    modelname = (
+        f"xnnpack_{modelname}" if xnnpack_extended_ops else f"xnnpack_dq_{modelname}"
+    )
 
     logging.info("Lowering model using following partitioner(s): ")
     for partitioner in partitioners:

examples/models/llama/tests/test_export_llama_lib.py

@@ -293,6 +293,50 @@ def test_has_expected_ops_and_op_counts(self):
         for op, _op_info in delegation_info.delegation_by_operator.items():
             self.assertTrue(op not in UNWANTED_OPS)
 
+    def test_xnnpack_extended_ops_defaults_on(self):
+        args = build_args_parser().parse_args([])
+        self.assertTrue(args.xnnpack_extended_ops)
+        self.assertTrue(LlmConfig.from_args(args).backend.xnnpack.extended_ops)
+
+    def test_no_xnnpack_extended_ops_opts_out(self):
+        args = build_args_parser().parse_args(["--no-xnnpack-extended-ops"])
+        self.assertFalse(args.xnnpack_extended_ops)
+        self.assertFalse(LlmConfig.from_args(args).backend.xnnpack.extended_ops)
+
+    def test_no_xnnpack_extended_ops_builds_dq_partitioner_only(self):
+        # Covers the multimethod path. _to_edge_and_lower_llama_xnnpack gates on
+        # the same flag and is exercised by the full-export tests in this file.
+        from executorch.backends.xnnpack.partition.config.xnnpack_config import (
+            ConfigPrecisionType,
+        )
+        from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+            XnnpackDynamicallyQuantizedPartitioner,
+        )
+
+        def build(argv):
+            args = build_args_parser().parse_args(["-X", *argv])
+            return export_llama_lib._get_xnnpack_partitioners(LlmConfig.from_args(args))
+
+        def precisions(partitioner):
+            return {
+                precision
+                for config in partitioner.target_partitioner_configs.values()
+                for precision in config.enabled_precision_types
+            }
+
+        dq_only = build(["--no-xnnpack-extended-ops"])
+        self.assertEqual(len(dq_only), 1)
+        self.assertIsInstance(dq_only[0], XnnpackDynamicallyQuantizedPartitioner)
+        self.assertEqual(precisions(dq_only[0]), {ConfigPrecisionType.DYNAMIC_QUANT})
+
+        # The default adds a second, general partitioner. It is not a subclass
+        # instance, and it enables precisions the dynamic-quant one cannot.
+        both = build([])
+        self.assertEqual(len(both), 2)
+        self.assertIsInstance(both[0], XnnpackDynamicallyQuantizedPartitioner)
+        self.assertNotIsInstance(both[1], XnnpackDynamicallyQuantizedPartitioner)
+        self.assertIn(ConfigPrecisionType.FP32, precisions(both[1]))
+
     def test_bf16_xnnpack_delegates_linears_when_enabled(self):
         parser = build_args_parser()
         args = parser.parse_args([])

extension/llm/export/config/llm_config.py

@@ -531,12 +531,13 @@ class XNNPackConfig:
     Attributes:
         enabled: :)
         extended_ops: Whether to match more types of ops to delegates to XNNPack.
+            On by default; set to False to restrict the delegate to DQLinear.
         enable_bf16: Whether to delegate BF16 ops to XNNPack. The target runtime
             must have hardware support for XNNPACK's BF16 kernels.
     """
 
     enabled: bool = False
-    extended_ops: bool = False
+    extended_ops: bool = True
     enable_bf16: bool = False