Update

[ghstack-poisoned]

Author: swolchok

backends/arm/operator_support/__init__.py

@@ -6,7 +6,6 @@
 # pyre-unsafe
 
 from . import (  # noqa
-    bitwise_support,
     convolution_support,
     pool_2d_support,
     reduce_sum_support,

backends/arm/operator_support/bitwise_support.py

@@ -11,13 +11,13 @@
 from typing import final, Optional, Sequence, Type
 
 import torch
-
 import torch.fx as fx
+
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm._passes.fuse_quantized_activation_pass import (
     FuseQuantizedActivationPass,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import Tosa_0_80, TosaSpecification
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.operator_support import any_chain, chain, OperatorSupportBase
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
@@ -90,6 +90,7 @@ def tosa_support_factory(
     if not tosa_spec.support_float():
         negative_checks.append(NeedsDecompositionCheck())
         negative_checks.append(CheckProperQuantization())
+        negative_checks.append(EthosU55NotSupported(tosa_spec))
     return chain(
         any_chain(
             BaseTOSASupportList(),
@@ -111,6 +112,9 @@ def is_node_supported(
         supported = node.op == "call_function" and node.target in [
             exir_ops.edge.aten.abs.default,
             exir_ops.edge.aten.add.Tensor,
+            exir_ops.edge.aten.bitwise_and.Tensor,
+            exir_ops.edge.aten.bitwise_or.Tensor,
+            exir_ops.edge.aten.bitwise_xor.Tensor,
             exir_ops.edge.aten.expand_copy.default,
             exir_ops.edge.aten.cat.default,
             exir_ops.edge.aten.clamp.default,
@@ -170,6 +174,31 @@ def is_node_supported(
         return supported
 
 
+class EthosU55NotSupported(OperatorSupportBase):
+    """
+    Certain operators are not supported on U55. These are listed in `unsupported` in
+    is_node_supported().
+    """
+
+    def __init__(self, tosa_spec: TosaSpecification):
+        self.tosa_spec = tosa_spec
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
+            unsupported_ops = [
+                exir_ops.edge.aten.bitwise_and.Tensor,
+                exir_ops.edge.aten.bitwise_or.Tensor,
+                exir_ops.edge.aten.bitwise_xor.Tensor,
+            ]
+
+            if node.target in unsupported_ops:
+                return False
+
+        return True
+
+
 class NeedsDecompositionCheck(OperatorSupportBase):
     """
     Targeted operators need to be decomposed prior to quantization in order to get a pair of q-dq-nodes surrounding

backends/arm/operator_support/tosa_supported_operators.py

@@ -9,8 +9,9 @@ python_binary(
     par_style = "xar",
     deps = [
         "//caffe2:torch",
-        "//executorch/exir:lib",
+        "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer",
         "//executorch/devtools:lib",
+        "//executorch/exir:lib",
         "//executorch/exir/backend/test:backend_with_compiler_demo",
         "//executorch/exir/backend/test:op_partitioner_demo",
         "//executorch/devtools/bundled_program/serialize:lib",

docs/TARGETS

@@ -1,15 +1,11 @@
-# Using Android prebuilt libraries (AAR)
+# Using Android prebuilt library (AAR)
 
-We provide two prebuilt Android libraries (AAR), `executorch.aar` for generic use case (image/audio processing) and `executorch_llama.aar` for LLAMA use case.
+We provide a prebuilt Android library (AAR), `executorch.aar` for both generic (image/audio processing) and LLAMA use case.
 
-## Contents of libraries
+## Contents of library
 - `executorch.aar`
   - [Java library](https://github.com/pytorch/executorch/tree/main/extension/android/src/main/java/org/pytorch/executorch)
-  - JNI contains the JNI binding for [NativePeer.java](https://github.com/pytorch/executorch/blob/main/extension/android/src/main/java/org/pytorch/executorch/NativePeer.java) and ExecuTorch native library, including core ExecuTorch runtime libraries, XNNPACK backend, Portable kernels, Optimized kernels, and Quantized kernels.
-    - Comes with two ABI variants, arm64-v8a and x86_64.
-- `executorch_llama.aar`
-  - [Java library](https://github.com/pytorch/executorch/tree/main/extension/android/src/main/java/org/pytorch/executorch) (Note: it contains the same Java classes as the previous Java, but it does not contain the JNI binding for generic Module/NativePeer Java code).
-  - JNI contains the JNI binding for [LlamaModule.java](https://github.com/pytorch/executorch/blob/main/extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java) and ExecuTorch native library, including core ExecuTorch runtime libraries, XNNPACK backend, Portable kernels, Optimized kernels, Quantized kernels, and LLAMA-specific Custom ops library.
+  - JNI contains the JNI binding for the corresponding Java code, and ExecuTorch native library, including core ExecuTorch runtime libraries, XNNPACK backend, Portable kernels, Optimized kernels, Quantized kernels, and LLAMA-specific Custom ops library.
     - Comes with two ABI variants, arm64-v8a and x86_64.
 
 ## Downloading AAR

docs/source/android-prebuilt-library.md

@@ -232,27 +232,29 @@ def forward(
         if self.apply_output:
             logits = self.output(h)
 
-        if self.output_prune_map is not None:
-            # expand to original size so that downstream applications can use the logits as-is.
-            if self.generate_full_logits:
-                # (1, seq_len, pruned_size) -> (1, seq_len, original_size)
-                expanded_logits = torch.full(
-                    [logits.shape[0], logits.shape[1], self.vocab_size],
-                    float("-inf"),
-                    device=logits.device,
-                    dtype=logits.dtype,
-                )
-                expanded_logits[:, :, list(self.output_prune_map.values())] = logits
-            else:
-                # (1, pruned_size) -> (1, original_size)
-                expanded_logits = torch.full(
-                    [logits.shape[0], self.vocab_size],
-                    float("-inf"),
-                    device=logits.device,
-                    dtype=logits.dtype,
-                )
-                expanded_logits[:, list(self.output_prune_map.values())] = logits
-            logits = expanded_logits
+            if self.output_prune_map is not None:
+                # expand to original size so that downstream applications can use the logits as-is.
+                if self.generate_full_logits:
+                    # (1, seq_len, pruned_size) -> (1, seq_len, original_size)
+                    expanded_logits = torch.full(
+                        [logits.shape[0], logits.shape[1], self.vocab_size],
+                        float("-inf"),
+                        device=logits.device,
+                        dtype=logits.dtype,
+                    )
+                    expanded_logits[:, :, list(self.output_prune_map.values())] = logits
+                else:
+                    # (1, pruned_size) -> (1, original_size)
+                    expanded_logits = torch.full(
+                        [logits.shape[0], self.vocab_size],
+                        float("-inf"),
+                        device=logits.device,
+                        dtype=logits.dtype,
+                    )
+                    expanded_logits[:, list(self.output_prune_map.values())] = logits
+                logits = expanded_logits
+        else:
+            logits = h
 
         if attn_options_update is not None:
             return logits, attn_options_update

examples/models/llama/llama_transformer.py

@@ -119,11 +119,10 @@ def quantize(  # noqa C901
         # Check for required args
         if group_size is None:
             raise Exception("For 8da4w quantization, group size must be specified.")
-        from torchao.quantization.quant_api import Int8DynActInt4WeightQuantizer
 
-        model = Int8DynActInt4WeightQuantizer(
-            precision=torch_dtype, groupsize=group_size
-        ).quantize(model)
+        from torchao.quantization import int8_dynamic_activation_int4_weight, quantize_
+
+        quantize_(model, int8_dynamic_activation_int4_weight(group_size=group_size))
 
         if verbose:
             print("quantized model:", model)
@@ -663,7 +662,7 @@ def convert_for_runtime(self) -> nn.Module:
     def quantized_model(self) -> nn.Module:
         model_updated_state_dict = self.create_quantized_state_dict(self.packed)
         self.convert_for_runtime()
-        self.mod.load_state_dict(model_updated_state_dict)
+        self.mod.load_state_dict(model_updated_state_dict, assign=True)
         return self.mod
 
 

examples/models/llama/source_transformation/quantize.py

@@ -8,7 +8,6 @@ def define_common_targets():
         ],
         exported_headers = ["flat_tensor_data_map.h"],
         deps = [
-            "//executorch/extension/flat_tensor/serialize:generated_headers",
             "//executorch/runtime/core:core",
             "//executorch/runtime/core:evalue",
             "//executorch/runtime/core:named_data_map",
@@ -17,6 +16,7 @@ def define_common_targets():
         ],
         exported_deps = [
             "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
+            "//executorch/extension/flat_tensor/serialize:generated_headers",
         ],
         visibility = [
             "//executorch/...",

extension/flat_tensor/targets.bzl

@@ -47,9 +47,6 @@ def define_common_targets(is_fbcode=False):
             deps = [
                 "//executorch/extension/data_loader:file_data_loader",
                 "//executorch/extension/flat_tensor:flat_tensor_data_map",
-                "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
-                "//executorch/extension/flat_tensor/serialize:generated_headers",
-                "//executorch/extension/flat_tensor/serialize:schema",
                 "//executorch/runtime/core:named_data_map",
                 "//executorch/runtime/core/exec_aten:lib",
             ],