From f7bb1c54d516fc7b322e4b58b9784b7bfaf550bf Mon Sep 17 00:00:00 2001 From: borys3kk Date: Thu, 22 Oct 2026 12:33:26 +0200 Subject: [PATCH 01/15] [XNNPACK] Lower upsample_nearest2d x2 and keep NHWC across scalar / low-rank broadcast operands Profiling a conv encoder-decoder (MI-GAN inpainting, fp16) showed the graph split into 7-8 XNNPACK partitions around portable upsample_nearest2d ops and ~25% of the time spent in NCHW<->NHWC transposes. upsample_nearest2d.vec with an exact x2 scale on a static 4D input is now partitioned and lowered as concatenate4 of the input along channels followed by depth_to_space_2d(block_size=2). XNNPACK has no nearest resize and its static_broadcast is not implemented at run time. This adds an XNNDepthToSpace2d node to the flatbuffer schemas and XNNCompiler. Static quantization is supported: the quantizer annotates the op with shared input/output qparams and the intermediate value carries them. Other scales, including an explicit scale_factors that is not exactly 2.0, stay portable. In ChannelsLastTaggedReshapePass, single-element operands of add/sub/mul/div/maximum/minimum (a lifted Python scalar such as the gain in x * 1.4142) no longer take part in the layout decision, and rank < 4 operands that broadcast over the trailing dims (a [H, W] map added to [N, C, H, W]) are reshaped to 4D with leading ones so they follow the NHWC operand instead of dragging the activation back to NCHW. Dequantized operands keep their rank, since a view between the dequantize and the op would carry no quant params. --- .../channels_last_tagged_reshape_pass.py | 88 ++++++++++++- backends/xnnpack/operators/__init__.py | 1 + .../operators/op_upsample_nearest2d.py | 124 ++++++++++++++++++ backends/xnnpack/partition/config/__init__.py | 2 + .../partition/config/generic_node_configs.py | 38 ++++++ .../xnnpack/quantizer/xnnpack_quantizer.py | 3 + .../quantizer/xnnpack_quantizer_utils.py | 35 +++++ backends/xnnpack/runtime/XNNCompiler.cpp | 34 +++++ .../xnnpack/serialization/runtime_schema.fbs | 8 ++ backends/xnnpack/serialization/schema.fbs | 8 ++ .../serialization/xnnpack_graph_schema.py | 9 ++ .../test/ops/test_upsample_nearest2d.py | 107 +++++++++++++++ .../test_channels_last_tagged_reshape.py | 115 ++++++++++++++++ 13 files changed, 565 insertions(+), 7 deletions(-) create mode 100644 backends/xnnpack/operators/op_upsample_nearest2d.py create mode 100644 backends/xnnpack/test/ops/test_upsample_nearest2d.py diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 5238c9689ee..15f156bc679 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -61,6 +61,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): memory_sensitive_ops_nhwc = { exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.aten.upsample_nearest2d.vec, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.amax.default, @@ -76,7 +77,9 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.linear.default, } - # Broadcasting binary ops whose operands must share one memory format. A + # Broadcasting binary ops whose operands must share one memory format. + # Single-element operands take no part in the layout decision, and rank < 4 + # operands are reshaped to 4D so they can follow an NHWC operand. A # dynamic-quant input_to_nhwc may blanket-replace a shared activation # (replace_all_uses_with), switching one operand of an already-processed # binary op to NHWC while its other (e.g. per-channel constant) operand @@ -86,6 +89,8 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.maximum.default, + exir_ops.edge.aten.minimum.default, } # Tag which is added to a node's meta to indicate that it uses NHWC format. @@ -173,6 +178,63 @@ def _is_nhwc_tensor(tensor: torch.Tensor) -> bool: def _is_nhwc(self, tensor: torch.Tensor) -> bool: return ChannelsLastTaggedReshapePass._is_nhwc_tensor(tensor) + @staticmethod + def is_single_element(node: torch.fx.Node) -> bool: + val = node.meta.get("val", None) + return isinstance(val, torch.Tensor) and val.numel() == 1 + + @staticmethod + def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: + val = node.meta.get("val", None) + return ( + isinstance(val, torch.Tensor) + and 0 < val.dim() < 4 + and val.numel() > 1 + and not ChannelsLastTaggedReshapePass.is_nhwc_node(node) + # A view between dequantize and the op would carry no qparams. + and not is_dequant(node) + ) + + def unsqueeze_to_4d( + self, + graph_module: torch.fx.GraphModule, + input_node: torch.fx.Node, + target_node: torch.fx.Node, + ) -> torch.fx.Node: + """ + Reshape a rank < 4 operand of `target_node` to 4D by prepending ones, + which is how broadcasting aligns it, so it can be converted to NHWC. + """ + val = input_node.meta["val"] + new_shape = [1] * (4 - val.dim()) + list(val.shape) + with graph_module.graph.inserting_before(target_node): + view = self.create_call_function_node( + graph_module=graph_module, + target=exir_ops.edge.aten.view_copy.default, + args=(input_node, new_shape), + ) + view.meta["val"] = val.reshape(new_shape) + # Left untagged on purpose: it is NCHW, and the view_copy visitor + # treats the presence of the tag key as "permute the shape". + target_node.replace_input_with(input_node, view) + return view + + def layout_inputs(self, node: torch.fx.Node) -> List[torch.fx.Node]: + """ + Inputs that take part in the layout decision of `node`. Single-element + operands of a broadcasting binary op (e.g. the scalar in `x * 1.4142`) + broadcast the same way in either layout, so they must not force an + NHWC activation back to NCHW. + """ + return [ + input_node + for input_node in node.all_input_nodes + if not ( + node.target in self.broadcast_binary_ops + and self.is_single_element(input_node) + ) + ] + def requires_nhwc_input(self, node: torch.fx.Node) -> bool: return node.target in self.memory_sensitive_ops_nhwc @@ -567,28 +629,40 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 else: # The node can have inputs in any format (but all must be the # same format) + layout_inputs = self.layout_inputs(node) is_or_isnt_nhwc_node = [ ChannelsLastTaggedReshapePass.is_nhwc_node(input_node) - for input_node in node.all_input_nodes + for input_node in layout_inputs ] - if all(is_or_isnt_nhwc_node): + if is_or_isnt_nhwc_node and all(is_or_isnt_nhwc_node): # All inputs are nhwc so this node's output is nhwc too self.mark_as_nhwc_node(node) elif any(is_or_isnt_nhwc_node): # There is at least one node of each format, so we need to # convert some of them to another format such that all # are the same. + # Reshape lower-rank operands (e.g. a [H, W] map added to + # [N, C, H, W]) to 4D so they can follow the NHWC operand. + if node.target in self.broadcast_binary_ops: + layout_inputs = [ + ( + self.unsqueeze_to_4d(graph_module, input_node, node) + if self.is_lower_rank_broadcast_operand(input_node) + else input_node + ) + for input_node in layout_inputs + ] # If all nodes can be converted to nhwc, then convert them # all to nhwc. Otherwise, convert all to nchw if all( self.can_be_converted_to_nhwc(input_node) - for input_node in node.all_input_nodes + for input_node in layout_inputs ): - for input_node in node.all_input_nodes: + for input_node in layout_inputs: self.input_to_nhwc(graph_module, input_node, node) self.mark_as_nhwc_node(node) else: - for input_node in node.all_input_nodes: + for input_node in layout_inputs: self.input_to_nchw(graph_module, input_node, node) graph_module.recompile() @@ -612,7 +686,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 or node.target not in ChannelsLastTaggedReshapePass.broadcast_binary_ops ): continue - input_nodes = node.all_input_nodes + input_nodes = self.layout_inputs(node) if len(input_nodes) != 2: continue layouts = [ diff --git a/backends/xnnpack/operators/__init__.py b/backends/xnnpack/operators/__init__.py index e06f337f9ee..4b7b93a8174 100644 --- a/backends/xnnpack/operators/__init__.py +++ b/backends/xnnpack/operators/__init__.py @@ -57,5 +57,6 @@ op_sub, op_tanh, op_to_copy, + op_upsample_nearest2d, op_view_copy, ) diff --git a/backends/xnnpack/operators/op_upsample_nearest2d.py b/backends/xnnpack/operators/op_upsample_nearest2d.py new file mode 100644 index 00000000000..6e96c24afab --- /dev/null +++ b/backends/xnnpack/operators/op_upsample_nearest2d.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Dict + +import torch +from executorch.backends.transforms import get_shape +from executorch.backends.xnnpack.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.xnnpack.operators.quant_params import QuantParams +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( + XNNConcatenate4, + XNNDepthToSpace2d, + XNNGraph, + XNNQuantizedTensorValue, + XNNTensorValue, + XNode, + XValue, +) +from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node +from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_INVALID_VALUE_ID + + +@register_node_visitor +class UpsampleNearest2dVisitor(NodeVisitor): + """ + XNNPACK has no nearest resize, so x2 upsampling is lowered as four copies + of the input concatenated along channels, [N, H, W, 4C], followed by + depth_to_space_2d(block_size=2). Neither op changes values, so the + intermediate carries the input's quantization parameters. + """ + + target = "aten.upsample_nearest2d.vec" + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + xnn_graph: XNNGraph, + vals_to_ids: Dict[torch.fx.Node, int], + debug_handle: int, + ) -> None: + input_node = get_input_node(node, 0) + input_quant_params = QuantParams.from_inputs(input_node, self._exported_program) + self.define_tensor( + input_node, + xnn_graph, + vals_to_ids, + convert_to_nhwc=True, + quant_params=input_quant_params, + ) + self.define_tensor( + node, + xnn_graph, + vals_to_ids, + convert_to_nhwc=True, + quant_params=QuantParams.from_outputs(node), + ) + input_id = vals_to_ids[input_node] + output_id = vals_to_ids[node] + + in_shape = get_shape(input_node) + out_shape = get_shape(node) + check_or_raise(len(in_shape) == 4, "upsample_nearest2d expects a 4D input") + n, c, h, w = in_shape + check_or_raise( + out_shape[-2] == 2 * h and out_shape[-1] == 2 * w, + f"only x2 nearest upsampling is supported, got {in_shape} -> {out_shape}", + ) + + concat_id = len(xnn_graph.xvalues) + concat_value = XNNTensorValue( + datatype=self.get_serialized_dtype(input_quant_params, node), + num_dims=4, + dims=[n, h, w, 4 * c], + external_id=XNN_INVALID_VALUE_ID, + constant_buffer_idx=0, + flags=0, + id_out=concat_id, + ) + xnn_graph.xvalues.append( + XValue(xvalue_union=concat_value) + if input_quant_params is None + else XValue( + xvalue_union=XNNQuantizedTensorValue( + tensor_value=concat_value, + quant_params=self.get_quant_params(input_quant_params, xnn_graph), + ) + ) + ) + + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNConcatenate4( + axis=3, + input1_id=input_id, + input2_id=input_id, + input3_id=input_id, + input4_id=input_id, + output_id=concat_id, + flags=0, + input5_id=XNN_INVALID_VALUE_ID, + ), + debug_handle=debug_handle, + ) + ) + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNDepthToSpace2d( + block_size=2, + input_id=concat_id, + output_id=output_id, + flags=0, + ), + debug_handle=debug_handle, + ) + ) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index 7775a674510..e0f4e560666 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -60,6 +60,7 @@ ToDimOrderCopyConfig, UnsqueezeCopyConfig, UpsampleBilinear2dConfig, + UpsampleNearest2dConfig, ViewCopyConfig, ) from executorch.backends.xnnpack.partition.config.node_configs import ( @@ -125,6 +126,7 @@ SubConfig, UnsqueezeCopyConfig, UpsampleBilinear2dConfig, + UpsampleNearest2dConfig, ViewCopyConfig, # Quant/Dequant Op Configs QuantizedPerTensorConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index d445b5ca8ac..0f80efcf138 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -503,6 +503,44 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]: return torch.ops.aten.upsample_bilinear2d.vec +class UpsampleNearest2dConfig(GenericNodePartitionerConfig): + target_name = "upsample_nearest2d.vec" + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + """ + Lowered as concatenate4 + depth_to_space_2d, which only expresses + exact x2 scaling of a static 4D tensor. + """ + if not self.check_common_constraints(node, ep): + return False + + if is_shape_dynamic(node): + why(node, reason="dynamic output sizes are not supported") + return False + + in_shape = get_input_node(node, 0).meta["val"].shape + out_shape = node.meta["val"].shape + if len(in_shape) != 4: + why(node, reason="only 4D inputs are supported") + return False + if out_shape[-2] != 2 * in_shape[-2] or out_shape[-1] != 2 * in_shape[-1]: + why(node, reason="only x2 nearest upsampling is supported") + return False + # PyTorch indexes with an explicit scale as given, so e.g. 2.004 + # doubles the size but picks different source pixels than x2. + scale_factors = node.args[2] if len(node.args) > 2 else None + if scale_factors is not None and any(s != 2.0 for s in scale_factors): + why(node, reason="only an exact x2 scale factor is supported") + return False + return True + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.upsample_nearest2d.vec + + class ExpConfig(GenericNodePartitionerConfig): target_name = "exp.default" diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer.py b/backends/xnnpack/quantizer/xnnpack_quantizer.py index 3c82a65ad71..cd2b4eb2f6c 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer.py @@ -272,6 +272,8 @@ class QuantPattern: CAT_TARGETS = {torch.ops.aten.cat.default} +UPSAMPLE_NEAREST2D_TARGETS = {torch.ops.aten.upsample_nearest2d.vec} + class XNNPACKQuantizer(Quantizer): supported_config_and_operators = _get_supported_config_and_operators() @@ -292,6 +294,7 @@ class XNNPACKQuantizer(Quantizer): QuantPattern("mul_relu", False, False, MUL_TARGETS), QuantPattern("mul", False, False, MUL_TARGETS), QuantPattern("cat", False, False, CAT_TARGETS), + QuantPattern("upsample_nearest2d", False, False, UPSAMPLE_NEAREST2D_TARGETS), ] def __init__(self) -> None: diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index a207d77674d..8755ea0c52b 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -727,6 +727,41 @@ def _annotate_adaptive_avg_pool2d( return annotated_partitions +@register_annotator("upsample_nearest2d") +def _annotate_upsample_nearest2d( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """Nearest upsampling only copies values, so its output shares the + input's quantization parameters.""" + annotated_partitions = [] + for node in gm.graph.nodes: + if node.target != torch.ops.aten.upsample_nearest2d.vec: + continue + if _is_annotated([node]) or (filter_fn and not filter_fn(node)): + continue + + input_act = node.args[0] + assert isinstance(input_act, Node) + if ( + Q_ANNOTATION_KEY not in input_act.meta + or not input_act.meta[Q_ANNOTATION_KEY]._annotated + or input_act.meta[Q_ANNOTATION_KEY].output_qspec is None + ): + input_act_qspec = get_input_act_qspec(quantization_config) + else: + input_act_qspec = SharedQuantizationSpec(input_act) + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map={input_act: input_act_qspec}, + output_qspec=SharedQuantizationSpec((input_act, node)), + _annotated=True, + ) + annotated_partitions.append([node]) + return annotated_partitions + + def _is_input_large_scalar(node: Node, gm: torch.fx.GraphModule): """Check if input is a large scalar value. So that we can skip quantization for the node since histc op (in HistogramObserver) only works for values up to certain upper bound diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 74a5eec4a91..c3b88bdc445 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -1642,6 +1642,39 @@ Error defineConcatenate5Node( return Error::Ok; } +/* +Defines serialized depth_to_space_2d node into the subgraph, +using the remapped ids to map the serialized ids, +to the new ids generated when defining the tensor value +*/ +Error defineDepthToSpace2dNode( + xnn_subgraph_t subgraph_ptr, + const std::unordered_map& remapped_ids, + const NodePtr node, + const fb_xnnpack::XNNGraph* graph) noexcept { + MAYBE_UNUSED(graph); + + auto graph_node = node->xnode_union_as_XNNDepthToSpace2d(); + + REMAP_ID(remapped_ids, graph_node->input_id(), d2s_input); + REMAP_ID(remapped_ids, graph_node->output_id(), d2s_output); + + xnn_status status = xnn_define_depth_to_space_2d( + subgraph_ptr, + graph_node->block_size(), + d2s_input, + d2s_output, + graph_node->flags()); + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Failed to create depth_to_space node %i with code: %s", + node->debug_handle(), + xnn_status_to_string(status)); + + return Error::Ok; +} + /* Defines serialized static_slice node into the subgraph, using the remapped ids to map the serialized ids, @@ -2052,6 +2085,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) { _DEFINE(Convert) _DEFINE(GlobalAvgPooling2d) _DEFINE(StaticReshape) + _DEFINE(DepthToSpace2d) _DEFINE(ArgMaxPooling2d) _DEFINE(Concatenate2) _DEFINE(Concatenate3) diff --git a/backends/xnnpack/serialization/runtime_schema.fbs b/backends/xnnpack/serialization/runtime_schema.fbs index 888384e032e..62143ab8c62 100644 --- a/backends/xnnpack/serialization/runtime_schema.fbs +++ b/backends/xnnpack/serialization/runtime_schema.fbs @@ -160,6 +160,7 @@ union XNodeUnion { XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, + XNNDepthToSpace2d, } union XValueUnion { @@ -298,6 +299,13 @@ table XNNStaticReshape { flags: uint; } +table XNNDepthToSpace2d { + block_size: uint; + input_id: uint; + output_id: uint; + flags: uint; +} + table XNNStaticSlice { num_dims:uint; offsets:[uint]; diff --git a/backends/xnnpack/serialization/schema.fbs b/backends/xnnpack/serialization/schema.fbs index 96e4256f078..a48f9bc4187 100644 --- a/backends/xnnpack/serialization/schema.fbs +++ b/backends/xnnpack/serialization/schema.fbs @@ -156,6 +156,7 @@ union XNodeUnion { XNNCopy: _XNNNode1x1, XNNCos: _XNNNode1x1, XNNApproxGelu: _XNNNode1x1, + XNNDepthToSpace2d, } union XValueUnion { @@ -294,6 +295,13 @@ table XNNStaticReshape { flags: uint; } +table XNNDepthToSpace2d { + block_size: uint; + input_id: uint; + output_id: uint; + flags: uint; +} + table XNNStaticSlice { num_dims:uint; offsets:[uint]; diff --git a/backends/xnnpack/serialization/xnnpack_graph_schema.py b/backends/xnnpack/serialization/xnnpack_graph_schema.py index 360034c1987..e6da4338dd0 100644 --- a/backends/xnnpack/serialization/xnnpack_graph_schema.py +++ b/backends/xnnpack/serialization/xnnpack_graph_schema.py @@ -202,6 +202,14 @@ class XNNStaticTranspose: flags: int +@dataclass +class XNNDepthToSpace2d: + block_size: int + input_id: int + output_id: int + flags: int + + @dataclass class XNNStaticSlice: num_dims: int @@ -427,6 +435,7 @@ class XNNScaledDotProductAttention: XNNCopy, XNNCos, XNNApproxGelu, + XNNDepthToSpace2d, ] diff --git a/backends/xnnpack/test/ops/test_upsample_nearest2d.py b/backends/xnnpack/test/ops/test_upsample_nearest2d.py new file mode 100644 index 00000000000..a4a105e6b09 --- /dev/null +++ b/backends/xnnpack/test/ops/test_upsample_nearest2d.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.xnnpack.test.tester import Tester + + +class TestUpsampleNearest2d(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class Nearest2x(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + + class Nearest2xBySize(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate( + x, size=(x.shape[2] * 2, x.shape[3] * 2), mode="nearest" + ) + + class Nearest3x(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.interpolate(x, scale_factor=3.0, mode="nearest") + + class NearestInexact2x(torch.nn.Module): + """Doubles the size, but PyTorch indexes with 1/2.004, so the x2 + duplication lowering would not match eager.""" + + def forward(self, x): + return torch.nn.functional.interpolate( + x, scale_factor=2.004, mode="nearest" + ) + + class ConvNearestConv(torch.nn.Module): + """The pattern that used to split XNNPACK partitions in conv decoders.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + + def forward(self, x): + x = self.conv1(x) + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv2(x) + + edge_op = "executorch_exir_dialects_edge__ops_aten_upsample_nearest2d_vec" + + def _test_lowered(self, module, inputs, quant=False): + tester = Tester(module, inputs) + if quant: + tester.quantize() + tester.export().to_edge_transform_and_lower().check_not([self.edge_op]) + if quant: + tester.check_not(["torch.ops.quantized_decomposed"]) + ( + tester.check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=2 if quant else 0) + ) + + def test_fp32_upsample_nearest2d_scale_factor(self): + self._test_lowered(self.Nearest2x(), (torch.randn(2, 3, 4, 5),)) + + def test_fp32_upsample_nearest2d_size(self): + self._test_lowered(self.Nearest2xBySize(), (torch.randn(1, 5, 7, 3),)) + + def test_fp16_upsample_nearest2d(self): + self._test_lowered( + self.Nearest2x(), (torch.randn(1, 6, 8, 8).to(torch.float16),) + ) + + def test_fp32_conv_upsample_conv_single_partition(self): + self._test_lowered(self.ConvNearestConv(), (torch.randn(1, 3, 16, 16),)) + + def test_qs8_conv_upsample_conv_single_partition(self): + self._test_lowered( + self.ConvNearestConv(), (torch.randn(1, 3, 16, 16),), quant=True + ) + + def test_fp32_upsample_nearest2d_x3_not_partitioned(self): + ( + Tester(self.Nearest3x(), (torch.randn(1, 3, 4, 4),)) + .export() + .to_edge_transform_and_lower() + .check_count({self.edge_op: 1}) + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + ) + + def test_fp32_upsample_nearest2d_inexact_scale_not_partitioned(self): + ( + Tester(self.NearestInexact2x(), (torch.randn(1, 3, 200, 200),)) + .export() + .to_edge_transform_and_lower() + .check_count({self.edge_op: 1}) + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index bfe2ce862c9..53bb958df06 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -290,6 +290,121 @@ def forward(self, x): x = torch.mean(x, (-1, -2), keepdim=True) return x + class ConvScalarMulConv(torch.nn.Module): + """conv -> x * scalar -> clamp -> conv: the scalar must not pull the + activation back to NCHW between the two convolutions.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + + def forward(self, x): + x = self.conv1(x) + x = torch.clamp(x * 1.4142, -256.0, 256.0) + return self.conv2(x) + + def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_stays_nhwc(self): + ( + Tester(self.ConvScalarMulConv().eval(), (torch.randn(1, 3, 8, 8),)) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC and output -> NCHW only + .check_count({self.to_copy_name: 2}) + .run_method_and_compare_outputs() + ) + + class ConvNoiseAddConv(torch.nn.Module): + """conv -> x + noise[H, W] -> conv: the rank-2 operand is reshaped to + [1, 1, H, W] and converted, instead of converting the activation.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8)) + + def forward(self, x): + x = self.conv1(x) + x = x + self.noise + return self.conv2(x) + + class ConvChannelScaleConv(torch.nn.Module): + """conv -> x * scale[C, 1, 1] -> conv: a per-channel operand is + reshaped to [1, C, 1, 1] and converted.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("scale", torch.randn(8, 1, 1)) + + def forward(self, x): + x = self.conv1(x) + x = x * self.scale + return self.conv2(x) + + class ConvNoiseSubConv(torch.nn.Module): + """conv -> noise[H, W] - x -> conv: the lower-rank operand is the + left argument.""" + + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) + self.register_buffer("noise", torch.randn(8, 8)) + + def forward(self, x): + x = self.conv1(x) + x = self.noise - x + return self.conv2(x) + + def _test_lower_rank_operand_stays_nhwc(self, module): + ( + Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + .export() + .to_edge() + .run_passes(self.PassStage) + # input -> NHWC, the 4D view of the operand -> NHWC, output -> NCHW + .check_count( + { + self.to_copy_name: 3, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + } + ) + .run_method_and_compare_outputs() + ) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): + self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseAddConv()) + + def test_fp32_channels_last_tagged_reshape_pass_per_channel_operand_stays_nhwc( + self, + ): + self._test_lower_rank_operand_stays_nhwc(self.ConvChannelScaleConv()) + + def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nhwc( + self, + ): + self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseSubConv()) + + def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): + # The dequantized operand is left at its own rank: a view between + # dequantize and the op would carry no quant params and fail at load. + for module in (self.ConvNoiseAddConv(), self.ConvChannelScaleConv()): + ( + Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) + .quantize() + .export() + .to_edge() + .run_passes(self.PassStage) + .check_not( + ["executorch_exir_dialects_edge__ops_aten_view_copy_default"] + ) + .run_method_and_compare_outputs() + ) + def test_fp32_channels_last_tagged_reshape_pass_conv_bn_hardtanh_mean_seq(self): # Copy #1 is for input to conv, nchw -> nhwc # Copy #2 is for conv to _native_batch_norm_legit_no_training, nhwc -> nchw From d6e8816d15989d3f77205aa3c22c2dcdc29ecb34 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Wed, 21 Oct 2026 15:27:22 +0200 Subject: [PATCH 02/15] Make delete_constant_placeholder edit the graph signature in place Backend passes keep a reference to the ExportedProgram they were built from (XNNPACKPass.exported_program, for one), while ExportedProgramPassManager runs them on a shallow copy of it and rebuilds the signature of that copy afterwards. delete_constant_placeholder replaced the original program's signature object, so the copy kept the old one and its verifier failed with Parameter is not in the state dict. create_constant_placeholder already inserts into the shared input_specs list in place; do the same for deletion. Tests cover a deletion made inside a pass run through the manager for every constant placeholder kind. Authored with Claude Code. --- ...test_create_delete_constant_placeholder.py | 58 +++++++++++++++++++ backends/transforms/utils.py | 15 ++--- 2 files changed, 63 insertions(+), 10 deletions(-) diff --git a/backends/transforms/test/test_create_delete_constant_placeholder.py b/backends/transforms/test/test_create_delete_constant_placeholder.py index 9444e6a4b63..4c4bc9f7822 100644 --- a/backends/transforms/test/test_create_delete_constant_placeholder.py +++ b/backends/transforms/test/test_create_delete_constant_placeholder.py @@ -10,6 +10,8 @@ ) from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from executorch.exir.pass_manager import ExportedProgramPassManager from torch.export import export from torch.export.graph_signature import InputKind @@ -204,3 +206,59 @@ def test_create_delete_non_persistent_buffer(): def test_create_delete_constant_tensor(): _test_create_delete(InputKind.CONSTANT_TENSOR) + + +def _test_delete_inside_a_pass(kind: InputKind, persistent_buffer: bool = None): + """Backend passes keep a reference to the program they were built from, + while ExportedProgramPassManager runs them on a shallow copy of it. The + deletion has to edit the shared signature in place, or the copy that the + manager verifies and returns still lists the deleted placeholder.""" + module = EmptyNetwork() + exported_program = export(module, args=module.test_data, strict=True) + exported_program = to_edge(exported_program).exported_program() + graph = exported_program.graph_module.graph + + const_name = "test_node" + input_node = list(graph.nodes)[0] + with graph.inserting_before(input_node): + const_node = create_constant_placeholder( + exp_program=exported_program, + graph=graph, + kind=kind, + name=const_name, + data=torch.ones(1), + persistent_buffer=persistent_buffer, + ) + assert len(exported_program.graph_signature.input_specs) == 2 + + def delete_unused_constant(graph_module: torch.fx.GraphModule) -> PassResult: + delete_constant_placeholder(exported_program, const_node) + return PassResult(graph_module, True) + + result = ExportedProgramPassManager([delete_unused_constant])(exported_program) + transformed = result.exported_program + signature = transformed.graph_signature + + assert len(signature.input_specs) == 1 + assert const_name not in signature.inputs_to_parameters + assert const_name not in signature.inputs_to_buffers + assert const_name not in signature.inputs_to_lifted_tensor_constants + assert len(transformed.state_dict) == 0 + assert len(transformed.constants) == 0 + assert transformed.module()(torch.zeros(1)) == 0 + + +def test_delete_inside_a_pass_parameter(): + _test_delete_inside_a_pass(InputKind.PARAMETER) + + +def test_delete_inside_a_pass_persistent_buffer(): + _test_delete_inside_a_pass(InputKind.BUFFER, True) + + +def test_delete_inside_a_pass_non_persistent_buffer(): + _test_delete_inside_a_pass(InputKind.BUFFER, False) + + +def test_delete_inside_a_pass_constant_tensor(): + _test_delete_inside_a_pass(InputKind.CONSTANT_TENSOR) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index d3c50c692b1..0ecfb99c8c0 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -266,16 +266,11 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod f"Cannot delete input node {node.name} since it is not a parameter, a buffer, nor a lifted tensor constant." ) - # Remove input from graph signature - input_specs = [ - spec - for spec in exp_program.graph_signature.input_specs - if spec.arg.name != node.name - ] - new_graph_signature = ExportGraphSignature( - input_specs, exp_program.graph_signature.output_specs - ) - exp_program._graph_signature = new_graph_signature + # Edit the spec list in place: ExportedProgramPassManager runs passes on a + # shallow copy of the program, which would not see a replaced signature. + input_specs = exp_program.graph_signature.input_specs + for spec in [spec for spec in input_specs if spec.arg.name == node.name]: + input_specs.remove(spec) # Remove node from graph node.graph.erase_node(node) From b929250ab985f4c71dea3962b37e2eadbf0ca340 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 20 Oct 2026 16:14:35 +0200 Subject: [PATCH 03/15] [XNNPACK] Lower depthwise transposed convs as polyphase depthwise convs + depth_to_space XNNPACK runs transposed convolutions through its subconvolution IGEMM path one group at a time. With one channel per group (a depthwise ConvTranspose2d, the natural way to write a StyleGAN upfirdn x2 upsampling filter) that is far slower than the depthwise conv microkernels: on MI-GAN the twelve 4x4 / stride-2 filters more than double the forward. DecomposeDepthwiseConvTransposePass rewrites such a node, when it is a clean x s upscale ((H-1)s - 2p + k + output_padding == sH, dilation 1, every phase's taps inside the input), into s*s stride-1 depthwise convs on the low-resolution input (phase (ay, ax) correlates with every s-th tap of the kernel, through a possibly asymmetric zero pad), a channel concatenation in (ay, ax) order and a block-major depth_to_space_2d. That last op is XNNPACK's native channel order (aten.pixel_shuffle is channel-major), so it is a new et_xnnpack custom op created only after partitioning, with a reference implementation for tracing and tests, a node visitor emitting XNNDepthToSpace2d, and an NHWC tag in the channels-last pass. The pass runs before FuseActivationPass, so an activation after the transposed conv is never fused into the node it removes. Every convolution is planned before any is rewritten, so deconvs feeding each other work, and the phase weights keep the original weight's input kind (parameter, buffer or constant). Deleting the folded weight relies on delete_constant_placeholder editing the graph signature in place. Authored with Claude Code. --- backends/xnnpack/BUCK | 11 + backends/xnnpack/_passes/BUCK | 1 + backends/xnnpack/_passes/__init__.py | 4 + .../channels_last_tagged_reshape_pass.py | 2 + .../decompose_depthwise_conv_transpose.py | 218 ++++++++++++++ backends/xnnpack/custom_ops_lib.py | 43 +++ backends/xnnpack/operators/BUCK | 1 + backends/xnnpack/operators/__init__.py | 1 + .../xnnpack/operators/op_depth_to_space.py | 70 +++++ ...test_decompose_depthwise_conv_transpose.py | 267 ++++++++++++++++++ 10 files changed, 618 insertions(+) create mode 100644 backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py create mode 100644 backends/xnnpack/custom_ops_lib.py create mode 100644 backends/xnnpack/operators/op_depth_to_space.py create mode 100644 backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py diff --git a/backends/xnnpack/BUCK b/backends/xnnpack/BUCK index 1e12548bc4c..f26b7fb6b0c 100644 --- a/backends/xnnpack/BUCK +++ b/backends/xnnpack/BUCK @@ -45,3 +45,14 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/xnnpack/utils:xnnpack_utils", ], ) + +fbcode_target(_kind = runtime.python_library, + name = "custom_ops_lib", + srcs = [ + "custom_ops_lib.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + ], +) diff --git a/backends/xnnpack/_passes/BUCK b/backends/xnnpack/_passes/BUCK index 2edf294fa4e..081b0ff31b2 100644 --- a/backends/xnnpack/_passes/BUCK +++ b/backends/xnnpack/_passes/BUCK @@ -11,6 +11,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/transforms:addmm_mm_to_linear", "//executorch/backends/transforms:remove_clone_ops", "//executorch/backends/transforms:lib", + "//executorch/backends/xnnpack:custom_ops_lib", "//executorch/backends/xnnpack/partition:partitioner_graphs", "//executorch/backends/xnnpack/serialization:xnnpack_schema", "//executorch/backends/xnnpack/utils:xnnpack_utils", diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index f55144395af..383734d466a 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -26,6 +26,9 @@ ) from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate +from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( + DecomposeDepthwiseConvTransposePass, +) from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass @@ -82,6 +85,7 @@ def __init__( ConstPropPass, FuseBatchNormPass, DecomposeBatchNorm, + DecomposeDepthwiseConvTransposePass, FuseActivationPass, DecomposeConcatenate, RemoveGetItemPass, diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 15f156bc679..31e2e71214b 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -9,6 +9,7 @@ from enum import Enum from typing import List, Optional, Tuple +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops import torch from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.utils.quant_utils import ( @@ -62,6 +63,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.upsample_bilinear2d.vec, exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.et_xnnpack.depth_to_space_2d.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.amax.default, diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py new file mode 100644 index 00000000000..875d7de5f11 --- /dev/null +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -0,0 +1,218 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import cast, List, Optional, Tuple + +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops +import torch +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + delete_constant_placeholder, + get_param_tensor, + is_param_node, +) +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult + + +class DecomposeDepthwiseConvTransposePass(XNNPACKPass): + """ + Lower a depthwise transposed conv (groups == in_channels == out_channels) + with stride s > 1 into s*s depthwise convs on the low-resolution input, a + channel concatenation and a block-major depth_to_space_2d. + + XNNPACK runs transposed convs through its subconvolution IGEMM path, one + group at a time; with one channel per group that is far slower than its + depthwise conv microkernels. Polyphase decomposition: output phase a of a + stride-s transposed conv is an ordinary stride-1 correlation of the input + with every s-th tap of the kernel, + + y[s*m + a] = sum_t x[m - t] * w[s*t + a + p], 0 <= s*t + a + p < k, + + per axis. The s*s phase outputs concatenated along channels in (ay, ax) + order are exactly the block-major input depth_to_space_2d expects. Each + phase reads the input through a (possibly asymmetric) zero pad. + + Only applied when the transposed conv is a clean x s upscale, i.e. + (H - 1)*s - 2p + k + output_padding == s*H on both axes with the same s, + dilation is 1 and every phase's taps stay inside the input; any other + transposed conv is left to XNNPACK's own node. + """ + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + # Plan every conv before rewriting any: the nodes created by + # _decompose have no meta["val"] until the retrace below. + plans = [] + for node in graph.nodes: + if ( + node.op == "call_function" + and node.target == exir_ops.edge.aten.convolution.default + ): + plan = self._plan(node) + if plan is not None: + plans.append((node, plan)) + for node, plan in plans: + self._decompose(graph, node, *plan) + + modified = bool(plans) + if modified: + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) + + # (t_min, t_max) per phase for one axis, None when the axis + # does not decompose. + @staticmethod + def _axis_phases( + size: int, kernel: int, stride: int, padding: int, output_padding: int + ) -> Optional[List[Tuple[int, int]]]: + if stride < 2: + return None + if (size - 1) * stride - 2 * padding + kernel + output_padding != size * stride: + return None + phases = [] + for a in range(stride): + # integer t with 0 <= stride*t + a + padding <= kernel - 1 + t_min = math.ceil(-(a + padding) / stride) + t_max = math.floor((kernel - 1 - a - padding) / stride) + if t_min > 0 or t_max < 0: + return None + phases.append((t_min, t_max)) + return phases + + def _plan(self, node: torch.fx.Node): + ( + x, + weight, + _bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) = node.args + if not transposed: + return None + if not is_param_node(self.exported_program, weight): + return None + x_shape = list(x.meta["val"].shape) + w_shape = list(weight.meta["val"].shape) + if len(x_shape) != 4 or len(w_shape) != 4: + return None + channels = x_shape[1] + # transposed conv weight is (in, out / groups, kh, kw) + if groups != channels or w_shape[0] != channels or w_shape[1] != 1: + return None + if any(d != 1 for d in dilation): + return None + # ATen broadcasts a single spatial value over both spatial dims. + stride, padding, output_padding = ( + p * 2 if len(p) == 1 else p + for p in cast(List[List[int]], [stride, padding, output_padding]) + ) + if stride[0] != stride[1]: + return None + phases_y = self._axis_phases( + x_shape[2], w_shape[2], stride[0], padding[0], output_padding[0] + ) + phases_x = self._axis_phases( + x_shape[3], w_shape[3], stride[1], padding[1], output_padding[1] + ) + if phases_y is None or phases_x is None: + return None + return stride[0], padding, phases_y, phases_x + + def _decompose( + self, + graph: torch.fx.Graph, + node: torch.fx.Node, + s: int, + padding: List[int], + phases_y: List[Tuple[int, int]], + phases_x: List[Tuple[int, int]], + ) -> None: + x, weight_node, bias_node = node.args[0], node.args[1], node.args[2] + groups = node.args[8] + weight = get_param_tensor(self.exported_program, weight_node) + assert weight is not None + weight_spec = next( + spec + for spec in self.exported_program.graph_signature.input_specs + if spec.arg.name == weight_node.name + ) + + phase_convs = [] + for ay, (ty_min, ty_max) in enumerate(phases_y): + for ax, (tx_min, tx_max) in enumerate(phases_x): + # taps of phase (ay, ax) as a stride-1 correlation kernel: + # W[q] = w[s * (t_max - q) + a + p] + iy = torch.tensor( + [ + s * (ty_max - q) + ay + padding[0] + for q in range(ty_max - ty_min + 1) + ] + ) + ix = torch.tensor( + [ + s * (tx_max - q) + ax + padding[1] + for q in range(tx_max - tx_min + 1) + ] + ) + phase_weight = weight[:, :, iy][:, :, :, ix].contiguous() + with graph.inserting_before(weight_node): + phase_weight_node = create_constant_placeholder( + exp_program=self.exported_program, + graph=graph, + kind=weight_spec.kind, + persistent_buffer=weight_spec.persistent, + name=f"{weight_node.name}_phase_{ay}_{ax}", + data=phase_weight, + ) + # (left, right, top, bottom) + pad = [tx_max, -tx_min, ty_max, -ty_min] + with graph.inserting_before(node): + conv_input = x + if any(pad): + conv_input = graph.create_node( + "call_function", + exir_ops.edge.aten.constant_pad_nd.default, + (x, pad, 0.0), + ) + phase_convs.append( + graph.create_node( + "call_function", + exir_ops.edge.aten.convolution.default, + ( + conv_input, + phase_weight_node, + bias_node, + [1, 1], + [0, 0], + [1, 1], + False, + [0, 0], + groups, + ), + ) + ) + + with graph.inserting_before(node): + cat = graph.create_node( + "call_function", exir_ops.edge.aten.cat.default, (phase_convs, 1) + ) + d2s = graph.create_node( + "call_function", + exir_ops.edge.et_xnnpack.depth_to_space_2d.default, + (cat, s), + ) + node.replace_all_uses_with(d2s) + graph.erase_node(node) + if len(weight_node.users) == 0: + delete_constant_placeholder(self.exported_program, weight_node) diff --git a/backends/xnnpack/custom_ops_lib.py b/backends/xnnpack/custom_ops_lib.py new file mode 100644 index 00000000000..5c393b54238 --- /dev/null +++ b/backends/xnnpack/custom_ops_lib.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Custom ops that XNNPACK passes introduce after partitioning. They never reach +the partitioner: a pass creates them inside an already lowered subgraph and a +node visitor serializes them, the reference implementation below only serves +tracing and tests. +""" + +import torch.library + +namespace = "et_xnnpack" +lib = torch.library.Library(namespace, "DEF") + +####################### +## depth_to_space_2d ## +####################### + + +def depth_to_space_2d_impl(x: torch.Tensor, block_size: int) -> torch.Tensor: + """ + Block-major depth to space, the channel order XNNPACK's (and TFLite's) + DEPTH_TO_SPACE uses: input channel (by * block_size + bx) * C + c lands at + output pixel (h * block_size + by, w * block_size + bx), channel c. + aten.pixel_shuffle is the channel-major variant, c * block_size**2 + by * + block_size + bx, and does not map onto XNNPACK's node directly. + """ + n, c_in, h, w = x.shape + bs = block_size + c = c_in // (bs * bs) + y = x.reshape(n, bs, bs, c, h, w) + y = y.permute(0, 3, 4, 1, 5, 2) # n, c, h, by, w, bx + return y.reshape(n, c, h * bs, w * bs) + + +name = "depth_to_space_2d" +lib.define(f"{name}(Tensor x, int block_size) -> Tensor") +lib.impl(name, depth_to_space_2d_impl, "CompositeExplicitAutograd") +depth_to_space_2d_op = torch.ops.et_xnnpack.depth_to_space_2d diff --git a/backends/xnnpack/operators/BUCK b/backends/xnnpack/operators/BUCK index fdc2ed1db49..a27c5d14026 100644 --- a/backends/xnnpack/operators/BUCK +++ b/backends/xnnpack/operators/BUCK @@ -8,6 +8,7 @@ fbcode_target(_kind = runtime.python_library, srcs = glob(["*.py"]), visibility = ["PUBLIC"], deps = [ + "//executorch/backends/xnnpack:custom_ops_lib", "//executorch/backends/xnnpack/utils:xnnpack_utils", "//executorch/exir:graph_module", "//executorch/exir/backend:backend_details", diff --git a/backends/xnnpack/operators/__init__.py b/backends/xnnpack/operators/__init__.py index 4b7b93a8174..3ee54158f7b 100644 --- a/backends/xnnpack/operators/__init__.py +++ b/backends/xnnpack/operators/__init__.py @@ -19,6 +19,7 @@ op_clone, op_conv2d, op_cos, + op_depth_to_space, op_div, op_dynamic_dequantize_ops, op_dynamic_quantize_ops, diff --git a/backends/xnnpack/operators/op_depth_to_space.py b/backends/xnnpack/operators/op_depth_to_space.py new file mode 100644 index 00000000000..1d913128515 --- /dev/null +++ b/backends/xnnpack/operators/op_depth_to_space.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import cast, Dict + +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops +import torch +from executorch.backends.transforms import get_shape +from executorch.backends.xnnpack.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( + XNNDepthToSpace2d, + XNNGraph, + XNode, +) +from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node + + +@register_node_visitor +class DepthToSpace2dVisitor(NodeVisitor): + """ + et_xnnpack.depth_to_space_2d, the block-major depth to space that XNNPACK + provides natively. Introduced by DecomposeDepthwiseConvTransposePass, see + custom_ops_lib.py for the channel order. + """ + + target = "et_xnnpack.depth_to_space_2d.default" + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + xnn_graph: XNNGraph, + vals_to_ids: Dict[torch.fx.Node, int], + debug_handle: int, + ) -> None: + input_node = get_input_node(node, 0) + self.define_tensor(input_node, xnn_graph, vals_to_ids, convert_to_nhwc=True) + self.define_tensor(node, xnn_graph, vals_to_ids, convert_to_nhwc=True) + + block_size = cast(int, node.args[1]) + in_shape = get_shape(input_node) + out_shape = get_shape(node) + check_or_raise( + len(in_shape) == 4 + and block_size >= 1 + and in_shape[1] % (block_size * block_size) == 0 + and out_shape[2] == in_shape[2] * block_size + and out_shape[3] == in_shape[3] * block_size, + f"depth_to_space_2d(block_size={block_size}) cannot map {in_shape} -> {out_shape}", + ) + + xnn_graph.xnodes.append( + XNode( + xnode_union=XNNDepthToSpace2d( + block_size=block_size, + input_id=vals_to_ids[input_node], + output_id=vals_to_ids[node], + flags=0, + ), + debug_handle=debug_handle, + ) + ) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py new file mode 100644 index 00000000000..b1ffc54eb04 --- /dev/null +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -0,0 +1,267 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math +import unittest + +import torch +from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( + DecomposeDepthwiseConvTransposePass, +) +from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir.dialects._ops import ops as exir_ops + + +class TestDecomposeDepthwiseConvTransposePass(unittest.TestCase): + PassStage = RunPasses([DecomposeDepthwiseConvTransposePass]) + + conv = "executorch_exir_dialects_edge__ops_aten_convolution_default" + pad = "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default" + cat = "executorch_exir_dialects_edge__ops_aten_cat_default" + d2s = "executorch_exir_dialects_edge__ops_et_xnnpack_depth_to_space_2d_default" + + def setUp(self): + torch._dynamo.reset() + + class DwDeconv(torch.nn.Module): + def __init__( + self, kernel, stride, padding, output_padding=0, bias=True, channels=8 + ): + super().__init__() + self.deconv = torch.nn.ConvTranspose2d( + channels, + channels, + kernel, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=channels, + bias=bias, + ) + with torch.no_grad(): + self.deconv.weight.normal_() + if bias: + self.deconv.bias.normal_() + + def forward(self, x): + return self.deconv(x) + + class QuantizedDwDeconv(torch.nn.Module): + # The XNNPACK quantizer never annotates grouped transposed convs, so + # build the dequantized weight by hand. + def __init__(self): + super().__init__() + self.register_buffer( + "weight", torch.randint(-128, 127, (8, 1, 4, 4), dtype=torch.int8) + ) + + def forward(self, x): + weight = torch.ops.quantized_decomposed.dequantize_per_tensor( + self.weight, 0.02, 0, -128, 127, torch.int8 + ) + return torch.nn.functional.conv_transpose2d( + x, weight, stride=2, padding=1, groups=8 + ) + + def _check_pass(self, module, inputs, counts, tol=1e-3, dynamic_shapes=None): + return ( + Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) + .export() + .to_edge() + .run_passes(self.PassStage) + .check_count(counts) + .run_method_and_compare_outputs(atol=tol, rtol=tol) + ) + + def _check_lowered(self, module, inputs, tol=1e-3, dynamic_shapes=None): + ( + Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not([self.conv]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(num_runs=3, atol=tol, rtol=tol) + ) + + def test_upfirdn_4x4_stride2(self): + # StyleGAN-style x2 upsampling filter: 4 phases of 2x2 taps, each with + # a one-sided pad. + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_no_bias_fp16(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1, bias=False).half() + inputs = (torch.randn(1, 8, 5, 7).half(),) + # fp16 accumulation order differs from the reference: 1 ulp at |y| ~ 10 + self._check_pass(module, inputs, {self.conv: 4, self.d2s: 1}, tol=1e-2) + self._check_lowered(module, inputs, tol=1e-2) + + def test_3x3_stride2_output_padding(self): + # k=3, p=1, output_padding=1 keeps the x2 scaling; the phases have 1 + # and 2 taps per axis. + module = self.DwDeconv(kernel=3, stride=2, padding=1, output_padding=1) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 4, self.cat: 1, self.d2s: 1}) + self._check_lowered(module, inputs) + + def test_2x2_stride2_no_padding(self): + # Every phase is a single tap, no pads at all. + module = self.DwDeconv(kernel=2, stride=2, padding=0) + inputs = (torch.randn(1, 8, 4, 6),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 0, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_stride3(self): + module = self.DwDeconv(kernel=3, stride=3, padding=0, channels=4) + inputs = (torch.randn(1, 4, 4, 5),) + self._check_pass(module, inputs, {self.conv: 9, self.cat: 1, self.d2s: 1}) + self._check_lowered(module, inputs) + + def test_conv_before_deconv(self): + # The phase pads read the conv's output rather than a graph input. + module = torch.nn.Sequential( + torch.nn.Conv2d(8, 8, 3, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + ) + inputs = (torch.randn(1, 8, 6, 10),) + tester = self._check_pass( + module, inputs, {self.conv: 5, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + for node in graph.nodes: + if node.target == exir_ops.edge.aten.constant_pad_nd.default: + self.assertEqual( + node.args[0].target, exir_ops.edge.aten.convolution.default + ) + self._check_lowered(module, inputs) + + def test_two_deconvs_in_a_row(self): + # The second deconv's input is the first one's depth_to_space_2d. + module = torch.nn.Sequential( + self.DwDeconv(kernel=4, stride=2, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + ) + inputs = (torch.randn(1, 8, 4, 6),) + self._check_pass( + module, inputs, {self.conv: 8, self.pad: 8, self.cat: 2, self.d2s: 2} + ) + self._check_lowered(module, inputs) + + def test_batch_gt_1(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(3, 8, 6, 10),) + self._check_pass( + module, inputs, {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1} + ) + self._check_lowered(module, inputs) + + def test_single_element_spatial_params(self): + class Deconv(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8, 1, 4, 4)) + + def forward(self, x): + return torch.ops.aten.convolution.default( + x, self.weight, None, [2], [1], [1], True, [0], 8 + ) + + inputs = (torch.randn(1, 8, 6, 10),) + self._check_pass(Deconv(), inputs, {self.conv: 4, self.pad: 4, self.d2s: 1}) + self._check_lowered(Deconv(), inputs) + + def test_deconv_relu(self): + # The activation follows depth_to_space_2d instead of fusing into the + # removed transposed conv. + module = torch.nn.Sequential( + self.DwDeconv(kernel=4, stride=2, padding=1), torch.nn.ReLU() + ) + inputs = (torch.randn(1, 8, 6, 10),) + self._check_lowered(module, inputs) + + def test_dynamic_shape(self): + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + dynamic_shapes = ( + { + 2: torch.export.Dim("height", min=2, max=16), + 3: torch.export.Dim("width", min=2, max=16), + }, + ) + self._check_pass( + module, + inputs, + {self.conv: 4, self.pad: 4, self.cat: 1, self.d2s: 1}, + dynamic_shapes=dynamic_shapes, + ) + self._check_lowered(module, inputs, dynamic_shapes=dynamic_shapes) + + def test_quantized_is_left_alone(self): + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(self.QuantizedDwDeconv(), inputs, {self.conv: 1, self.d2s: 0}) + + def test_not_a_clean_upscale_is_left_alone(self): + # k=4, p=0: output is 2H + 2, not a x2 upscale. + module = self.DwDeconv(kernel=4, stride=2, padding=0) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 1, self.d2s: 0}) + + def test_grouped_but_not_depthwise_is_left_alone(self): + module = torch.nn.ConvTranspose2d(8, 8, 4, stride=2, padding=1, groups=4) + inputs = (torch.randn(1, 8, 6, 6),) + self._check_pass(module, inputs, {self.conv: 1, self.d2s: 0}) + + def test_phase_math_matches_conv_transpose2d(self): + # The decomposition in plain torch, independent of the lowering. + import executorch.backends.xnnpack.custom_ops_lib as lib + + for k, s, p, o in [ + (4, 2, 1, 0), + (3, 2, 1, 1), + (2, 2, 0, 0), + (3, 3, 0, 0), + (5, 2, 2, 1), + ]: + c, h, w = 3, 5, 6 + x = torch.randn(1, c, h, w) + weight = torch.randn(c, 1, k, k) + ref = torch.nn.functional.conv_transpose2d( + x, weight, stride=s, padding=p, output_padding=o, groups=c + ) + phases = [] + for ay in range(s): + for ax in range(s): + taps = [] + for a in (ay, ax): + t_min = math.ceil(-(a + p) / s) + t_max = math.floor((k - 1 - a - p) / s) + taps.append((t_min, t_max)) + (ty_min, ty_max), (tx_min, tx_max) = taps + iy = torch.tensor( + [s * (ty_max - q) + ay + p for q in range(ty_max - ty_min + 1)] + ) + ix = torch.tensor( + [s * (tx_max - q) + ax + p for q in range(tx_max - tx_min + 1)] + ) + wp = weight[:, :, iy][:, :, :, ix] + xp = torch.nn.functional.pad(x, (tx_max, -tx_min, ty_max, -ty_min)) + phases.append(torch.nn.functional.conv2d(xp, wp, groups=c)) + out = lib.depth_to_space_2d_op(torch.cat(phases, 1), s) + self.assertEqual(out.shape, ref.shape, (k, s, p, o)) + self.assertTrue(torch.allclose(out, ref, atol=1e-5), (k, s, p, o)) From 743c690f953855f56bfb3ff4a56de7e33b948c9f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 26 Oct 2026 13:21:24 +0100 Subject: [PATCH 04/15] Remove the deleted spec with a slice assignment --- backends/transforms/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 0ecfb99c8c0..fe36ff9ca69 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -269,8 +269,7 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod # Edit the spec list in place: ExportedProgramPassManager runs passes on a # shallow copy of the program, which would not see a replaced signature. input_specs = exp_program.graph_signature.input_specs - for spec in [spec for spec in input_specs if spec.arg.name == node.name]: - input_specs.remove(spec) + input_specs[:] = [spec for spec in input_specs if spec.arg.name != node.name] # Remove node from graph node.graph.erase_node(node) From 1c8a70d5a5d7a1d4950d7701abfbf22020436211 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 27 Oct 2026 08:48:53 +0100 Subject: [PATCH 05/15] Use a Software Mansion copyright header on the new files --- backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py | 2 +- backends/xnnpack/custom_ops_lib.py | 2 +- backends/xnnpack/operators/op_depth_to_space.py | 2 +- backends/xnnpack/operators/op_upsample_nearest2d.py | 2 +- backends/xnnpack/test/ops/test_upsample_nearest2d.py | 2 +- .../test/passes/test_decompose_depthwise_conv_transpose.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 875d7de5f11..94fac0a506c 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/custom_ops_lib.py b/backends/xnnpack/custom_ops_lib.py index 5c393b54238..290adfa6d84 100644 --- a/backends/xnnpack/custom_ops_lib.py +++ b/backends/xnnpack/custom_ops_lib.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/operators/op_depth_to_space.py b/backends/xnnpack/operators/op_depth_to_space.py index 1d913128515..0c7a2ae2baa 100644 --- a/backends/xnnpack/operators/op_depth_to_space.py +++ b/backends/xnnpack/operators/op_depth_to_space.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/operators/op_upsample_nearest2d.py b/backends/xnnpack/operators/op_upsample_nearest2d.py index 6e96c24afab..ec408a65e74 100644 --- a/backends/xnnpack/operators/op_upsample_nearest2d.py +++ b/backends/xnnpack/operators/op_upsample_nearest2d.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/test/ops/test_upsample_nearest2d.py b/backends/xnnpack/test/ops/test_upsample_nearest2d.py index a4a105e6b09..e70c4d4af72 100644 --- a/backends/xnnpack/test/ops/test_upsample_nearest2d.py +++ b/backends/xnnpack/test/ops/test_upsample_nearest2d.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index b1ffc54eb04..3cc2c3d8aec 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026 Software Mansion # All rights reserved. # # This source code is licensed under the BSD-style license found in the From 445b1d50f5e87319a823f017d41d4ad5f330372f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Tue, 27 Oct 2026 09:07:08 +0100 Subject: [PATCH 06/15] Drop the channels-last broadcast operand changes that moved out of #24113 --- .../channels_last_tagged_reshape_pass.py | 87 ++----------- .../test_channels_last_tagged_reshape.py | 115 ------------------ 2 files changed, 7 insertions(+), 195 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index 31e2e71214b..656d1214627 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -79,9 +79,7 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.linear.default, } - # Broadcasting binary ops whose operands must share one memory format. - # Single-element operands take no part in the layout decision, and rank < 4 - # operands are reshaped to 4D so they can follow an NHWC operand. A + # Broadcasting binary ops whose operands must share one memory format. A # dynamic-quant input_to_nhwc may blanket-replace a shared activation # (replace_all_uses_with), switching one operand of an already-processed # binary op to NHWC while its other (e.g. per-channel constant) operand @@ -91,8 +89,6 @@ class ChannelsLastTaggedReshapePass(XNNPACKPass): exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.div.Tensor, - exir_ops.edge.aten.maximum.default, - exir_ops.edge.aten.minimum.default, } # Tag which is added to a node's meta to indicate that it uses NHWC format. @@ -180,63 +176,6 @@ def _is_nhwc_tensor(tensor: torch.Tensor) -> bool: def _is_nhwc(self, tensor: torch.Tensor) -> bool: return ChannelsLastTaggedReshapePass._is_nhwc_tensor(tensor) - @staticmethod - def is_single_element(node: torch.fx.Node) -> bool: - val = node.meta.get("val", None) - return isinstance(val, torch.Tensor) and val.numel() == 1 - - @staticmethod - def is_lower_rank_broadcast_operand(node: torch.fx.Node) -> bool: - val = node.meta.get("val", None) - return ( - isinstance(val, torch.Tensor) - and 0 < val.dim() < 4 - and val.numel() > 1 - and not ChannelsLastTaggedReshapePass.is_nhwc_node(node) - # A view between dequantize and the op would carry no qparams. - and not is_dequant(node) - ) - - def unsqueeze_to_4d( - self, - graph_module: torch.fx.GraphModule, - input_node: torch.fx.Node, - target_node: torch.fx.Node, - ) -> torch.fx.Node: - """ - Reshape a rank < 4 operand of `target_node` to 4D by prepending ones, - which is how broadcasting aligns it, so it can be converted to NHWC. - """ - val = input_node.meta["val"] - new_shape = [1] * (4 - val.dim()) + list(val.shape) - with graph_module.graph.inserting_before(target_node): - view = self.create_call_function_node( - graph_module=graph_module, - target=exir_ops.edge.aten.view_copy.default, - args=(input_node, new_shape), - ) - view.meta["val"] = val.reshape(new_shape) - # Left untagged on purpose: it is NCHW, and the view_copy visitor - # treats the presence of the tag key as "permute the shape". - target_node.replace_input_with(input_node, view) - return view - - def layout_inputs(self, node: torch.fx.Node) -> List[torch.fx.Node]: - """ - Inputs that take part in the layout decision of `node`. Single-element - operands of a broadcasting binary op (e.g. the scalar in `x * 1.4142`) - broadcast the same way in either layout, so they must not force an - NHWC activation back to NCHW. - """ - return [ - input_node - for input_node in node.all_input_nodes - if not ( - node.target in self.broadcast_binary_ops - and self.is_single_element(input_node) - ) - ] - def requires_nhwc_input(self, node: torch.fx.Node) -> bool: return node.target in self.memory_sensitive_ops_nhwc @@ -631,40 +570,28 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 else: # The node can have inputs in any format (but all must be the # same format) - layout_inputs = self.layout_inputs(node) is_or_isnt_nhwc_node = [ ChannelsLastTaggedReshapePass.is_nhwc_node(input_node) - for input_node in layout_inputs + for input_node in node.all_input_nodes ] - if is_or_isnt_nhwc_node and all(is_or_isnt_nhwc_node): + if all(is_or_isnt_nhwc_node): # All inputs are nhwc so this node's output is nhwc too self.mark_as_nhwc_node(node) elif any(is_or_isnt_nhwc_node): # There is at least one node of each format, so we need to # convert some of them to another format such that all # are the same. - # Reshape lower-rank operands (e.g. a [H, W] map added to - # [N, C, H, W]) to 4D so they can follow the NHWC operand. - if node.target in self.broadcast_binary_ops: - layout_inputs = [ - ( - self.unsqueeze_to_4d(graph_module, input_node, node) - if self.is_lower_rank_broadcast_operand(input_node) - else input_node - ) - for input_node in layout_inputs - ] # If all nodes can be converted to nhwc, then convert them # all to nhwc. Otherwise, convert all to nchw if all( self.can_be_converted_to_nhwc(input_node) - for input_node in layout_inputs + for input_node in node.all_input_nodes ): - for input_node in layout_inputs: + for input_node in node.all_input_nodes: self.input_to_nhwc(graph_module, input_node, node) self.mark_as_nhwc_node(node) else: - for input_node in layout_inputs: + for input_node in node.all_input_nodes: self.input_to_nchw(graph_module, input_node, node) graph_module.recompile() @@ -688,7 +615,7 @@ def call(self, graph_module: torch.fx.GraphModule): # noqa: C901 or node.target not in ChannelsLastTaggedReshapePass.broadcast_binary_ops ): continue - input_nodes = self.layout_inputs(node) + input_nodes = node.all_input_nodes if len(input_nodes) != 2: continue layouts = [ diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index 53bb958df06..bfe2ce862c9 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -290,121 +290,6 @@ def forward(self, x): x = torch.mean(x, (-1, -2), keepdim=True) return x - class ConvScalarMulConv(torch.nn.Module): - """conv -> x * scalar -> clamp -> conv: the scalar must not pull the - activation back to NCHW between the two convolutions.""" - - def __init__(self): - super().__init__() - self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - - def forward(self, x): - x = self.conv1(x) - x = torch.clamp(x * 1.4142, -256.0, 256.0) - return self.conv2(x) - - def test_fp32_channels_last_tagged_reshape_pass_scalar_operand_stays_nhwc(self): - ( - Tester(self.ConvScalarMulConv().eval(), (torch.randn(1, 3, 8, 8),)) - .export() - .to_edge() - .run_passes(self.PassStage) - # input -> NHWC and output -> NCHW only - .check_count({self.to_copy_name: 2}) - .run_method_and_compare_outputs() - ) - - class ConvNoiseAddConv(torch.nn.Module): - """conv -> x + noise[H, W] -> conv: the rank-2 operand is reshaped to - [1, 1, H, W] and converted, instead of converting the activation.""" - - def __init__(self): - super().__init__() - self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8)) - - def forward(self, x): - x = self.conv1(x) - x = x + self.noise - return self.conv2(x) - - class ConvChannelScaleConv(torch.nn.Module): - """conv -> x * scale[C, 1, 1] -> conv: a per-channel operand is - reshaped to [1, C, 1, 1] and converted.""" - - def __init__(self): - super().__init__() - self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("scale", torch.randn(8, 1, 1)) - - def forward(self, x): - x = self.conv1(x) - x = x * self.scale - return self.conv2(x) - - class ConvNoiseSubConv(torch.nn.Module): - """conv -> noise[H, W] - x -> conv: the lower-rank operand is the - left argument.""" - - def __init__(self): - super().__init__() - self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = torch.nn.Conv2d(8, 4, 3, padding=1) - self.register_buffer("noise", torch.randn(8, 8)) - - def forward(self, x): - x = self.conv1(x) - x = self.noise - x - return self.conv2(x) - - def _test_lower_rank_operand_stays_nhwc(self, module): - ( - Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) - .export() - .to_edge() - .run_passes(self.PassStage) - # input -> NHWC, the 4D view of the operand -> NHWC, output -> NCHW - .check_count( - { - self.to_copy_name: 3, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - } - ) - .run_method_and_compare_outputs() - ) - - def test_fp32_channels_last_tagged_reshape_pass_lower_rank_operand_stays_nhwc(self): - self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseAddConv()) - - def test_fp32_channels_last_tagged_reshape_pass_per_channel_operand_stays_nhwc( - self, - ): - self._test_lower_rank_operand_stays_nhwc(self.ConvChannelScaleConv()) - - def test_fp32_channels_last_tagged_reshape_pass_lower_rank_left_operand_stays_nhwc( - self, - ): - self._test_lower_rank_operand_stays_nhwc(self.ConvNoiseSubConv()) - - def test_qs8_channels_last_tagged_reshape_pass_lower_rank_operand(self): - # The dequantized operand is left at its own rank: a view between - # dequantize and the op would carry no quant params and fail at load. - for module in (self.ConvNoiseAddConv(), self.ConvChannelScaleConv()): - ( - Tester(module.eval(), (torch.randn(1, 3, 8, 8),)) - .quantize() - .export() - .to_edge() - .run_passes(self.PassStage) - .check_not( - ["executorch_exir_dialects_edge__ops_aten_view_copy_default"] - ) - .run_method_and_compare_outputs() - ) - def test_fp32_channels_last_tagged_reshape_pass_conv_bn_hardtanh_mean_seq(self): # Copy #1 is for input to conv, nchw -> nhwc # Copy #2 is for conv to _native_batch_norm_legit_no_training, nhwc -> nchw From c22288b72ef0c73a8bc76efc5cc35a2a41d0e220 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Thu, 29 Oct 2026 12:16:52 +0100 Subject: [PATCH 07/15] Key the phase weight placeholders by stride and padding as well create_constant_placeholder returns an existing placeholder of the same name, so two depthwise deconvs sharing one weight but using different stride or padding ended up with the first one's phase kernels. Include both in the name and add tests for a shared weight with different and with identical parameters. --- .../decompose_depthwise_conv_transpose.py | 9 ++- ...test_decompose_depthwise_conv_transpose.py | 57 +++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 94fac0a506c..986615932fe 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -166,13 +166,20 @@ def _decompose( ] ) phase_weight = weight[:, :, iy][:, :, :, ix].contiguous() + # The phase taps depend on s and padding too, so a weight + # shared by differently configured deconvs must not reuse + # the placeholders create_constant_placeholder finds by name. + phase_name = ( + f"{weight_node.name}_s{s}_p{padding[0]}x{padding[1]}" + f"_phase_{ay}_{ax}" + ) with graph.inserting_before(weight_node): phase_weight_node = create_constant_placeholder( exp_program=self.exported_program, graph=graph, kind=weight_spec.kind, persistent_buffer=weight_spec.persistent, - name=f"{weight_node.name}_phase_{ay}_{ax}", + name=phase_name, data=phase_weight, ) # (left, right, top, bottom) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index 3cc2c3d8aec..8ff95e4d22c 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -163,6 +163,63 @@ def test_two_deconvs_in_a_row(self): ) self._check_lowered(module, inputs) + class SharedWeightDeconvs(torch.nn.Module): + def __init__(self, kernel, stride, paddings, output_paddings): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8, 1, kernel, kernel)) + self.stride = stride + self.configs = list(zip(paddings, output_paddings)) + + def forward(self, x): + return sum( + torch.nn.functional.conv_transpose2d( + x, + self.weight, + stride=self.stride, + padding=padding, + output_padding=output_padding, + groups=8, + ) + for padding, output_padding in self.configs + ) + + def _count_phase_placeholders(self, tester): + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + return sum( + 1 + for node in graph.nodes + if node.op == "placeholder" and "_phase_" in node.name + ) + + def test_shared_weight_different_padding(self): + # Both deconvs are clean x3 upscales of the same weight, but their + # phase taps differ, so they must not share the phase placeholders. + module = self.SharedWeightDeconvs( + kernel=9, stride=3, paddings=(3, 4), output_paddings=(0, 2) + ) + inputs = (torch.randn(1, 8, 5, 6),) + tester = self._check_pass( + module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} + ) + self.assertEqual(self._count_phase_placeholders(tester), 18) + self._check_lowered(module, inputs) + + def test_shared_weight_same_params(self): + # Identical deconvs on one weight share one set of phase placeholders. + module = self.SharedWeightDeconvs( + kernel=9, stride=3, paddings=(3, 3), output_paddings=(0, 0) + ) + inputs = (torch.randn(1, 8, 5, 6),) + tester = self._check_pass( + module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} + ) + self.assertEqual(self._count_phase_placeholders(tester), 9) + self._check_lowered(module, inputs) + def test_batch_gt_1(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) inputs = (torch.randn(3, 8, 6, 10),) From 850c9763dee4d7952ef5b81bb09457c1783a5f97 Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 2 Nov 2026 12:49:19 +0100 Subject: [PATCH 08/15] Keep the external-constant tag on the phase weights create_constant_placeholder only sets meta["val"], so the phase weights lost the delegate_constant_tag of the weight they replace and landed in the .pte instead of the .ptd file. --- .../xnnpack/_passes/decompose_depthwise_conv_transpose.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py index 986615932fe..c41e0be5eeb 100644 --- a/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/_passes/decompose_depthwise_conv_transpose.py @@ -182,6 +182,11 @@ def _decompose( name=phase_name, data=phase_weight, ) + # create_constant_placeholder only sets meta["val"]; keep the + # external-constant tag so the phase weights follow the + # original weight into its .ptd file. + if "custom" in weight_node.meta: + phase_weight_node.meta["custom"] = dict(weight_node.meta["custom"]) # (left, right, top, bottom) pad = [tx_max, -tx_min, ty_max, -ty_min] with graph.inserting_before(node): From 4ffc1bad8f1d1ac1f2b942af15750bbb03ee577f Mon Sep 17 00:00:00 2001 From: borys3kk Date: Mon, 2 Nov 2026 12:49:20 +0100 Subject: [PATCH 09/15] Test that the pass runs during lowering, its order against FuseActivationPass, and external constants XNNPACK delegates an fp32 depthwise transposed conv on its own, so the lowered tests passed whether or not the decomposition happened. _check_lowered now counts the _decompose calls and checks the serialized graph for depth_to_space_2d and no transposed conv. test_deconv_relu runs the default pass list and checks that the relu reads the depth_to_space_2d output and that the phase pads fold. A new test lowers with external constants and expects every constant in model.ptd. --- ...test_decompose_depthwise_conv_transpose.py | 178 ++++++++++++++++-- 1 file changed, 166 insertions(+), 12 deletions(-) diff --git a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py index 8ff95e4d22c..5bba9cc4479 100644 --- a/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py +++ b/backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py @@ -4,16 +4,39 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import math +import os +import tempfile import unittest +from collections import Counter +from unittest import mock +import executorch.backends.test.harness.stages as BaseStages import torch from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack._passes import XNNPACKPassManager from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( DecomposeDepthwiseConvTransposePass, ) +from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass +from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import ( + FusePadIntoConvPass, +) +from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import ( + _get_schema_bytes, + XNNHeader, +) from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir._serialize import _deserialize_pte_binary +from executorch.exir._serialize._flatbuffer import _flatc_decompile from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.external_constants_pass import ( + delegate_external_constants_pass_unlifted, +) +from executorch.extension.flat_tensor.serialize.serialize import ( + _deserialize_to_flat_tensor, +) class TestDecomposeDepthwiseConvTransposePass(unittest.TestCase): @@ -77,17 +100,68 @@ def _check_pass(self, module, inputs, counts, tol=1e-3, dynamic_shapes=None): .run_method_and_compare_outputs(atol=tol, rtol=tol) ) - def _check_lowered(self, module, inputs, tol=1e-3, dynamic_shapes=None): + @staticmethod + def _xnnpack_node_types(tester) -> Counter: + """Node types in the serialized XNNPACK graph of the single delegate + produced by to_edge_transform_and_lower.""" + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + (delegate,) = [ + node + for node in graph_module.graph.nodes + if node.target == torch.ops.higher_order.executorch_call_delegate + ] + payload = getattr(graph_module, delegate.args[0].target).processed_bytes + header = XNNHeader.from_bytes(payload[: XNNHeader.EXPECTED_LENGTH]) + flatbuffer = payload[ + header.flatbuffer_offset : header.flatbuffer_offset + header.flatbuffer_size + ] + with tempfile.TemporaryDirectory() as d: + schema_path = os.path.join(d, "schema.fbs") + with open(schema_path, "wb") as f: + f.write(_get_schema_bytes()) + bin_path = os.path.join(d, "graph.bin") + with open(bin_path, "wb") as f: + f.write(flatbuffer) + _flatc_decompile(d, schema_path, bin_path) + with open(os.path.join(d, "graph.json")) as f: + # flatc writes unbounded clamp limits as a bare inf + text = ( + f.read().replace("-inf", "-Infinity").replace(" inf", " Infinity") + ) + graph = json.loads(text) + return Counter(node["xnode_union_type"] for node in graph["xnodes"]) + + def _check_lowered( + self, module, inputs, tol=1e-3, dynamic_shapes=None, num_deconvs=1 + ): + # XNNPACK delegates an fp32 depthwise transposed conv on its own, so + # one delegate and matching outputs do not prove the pass ran during + # lowering: count its calls and check the serialized graph. + tester = Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes).export() + with mock.patch.object( + DecomposeDepthwiseConvTransposePass, + "_decompose", + autospec=True, + side_effect=DecomposeDepthwiseConvTransposePass._decompose, + ) as decompose: + tester.to_edge_transform_and_lower() + self.assertEqual(decompose.call_count, num_deconvs) + tester.check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ).check_not([self.conv]) + node_types = self._xnnpack_node_types(tester) + self.assertEqual(node_types["XNNConvTranspose2d"], 0, node_types) + self.assertEqual(node_types["XNNDepthToSpace2d"], num_deconvs, node_types) ( - Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) - .export() - .to_edge_transform_and_lower() - .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) - .check_not([self.conv]) - .to_executorch() + tester.to_executorch() .serialize() .run_method_and_compare_outputs(num_runs=3, atol=tol, rtol=tol) ) + return tester def test_upfirdn_4x4_stride2(self): # StyleGAN-style x2 upsampling filter: 4 phases of 2x2 taps, each with @@ -161,7 +235,7 @@ def test_two_deconvs_in_a_row(self): self._check_pass( module, inputs, {self.conv: 8, self.pad: 8, self.cat: 2, self.d2s: 2} ) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) class SharedWeightDeconvs(torch.nn.Module): def __init__(self, kernel, stride, paddings, output_paddings): @@ -206,7 +280,7 @@ def test_shared_weight_different_padding(self): module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} ) self.assertEqual(self._count_phase_placeholders(tester), 18) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) def test_shared_weight_same_params(self): # Identical deconvs on one weight share one set of phase placeholders. @@ -218,7 +292,7 @@ def test_shared_weight_same_params(self): module, inputs, {self.conv: 18, self.cat: 2, self.d2s: 2} ) self.assertEqual(self._count_phase_placeholders(tester), 9) - self._check_lowered(module, inputs) + self._check_lowered(module, inputs, num_deconvs=2) def test_batch_gt_1(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) @@ -245,13 +319,93 @@ def forward(self, x): def test_deconv_relu(self): # The activation follows depth_to_space_2d instead of fusing into the - # removed transposed conv. + # removed transposed conv, which relies on the pass running before + # FuseActivationPass in the default pass list. The phase pads read the + # NHWC output of the leading conv and fold into the phase convs through + # FusePadIntoConvPass at the end of that list. module = torch.nn.Sequential( - self.DwDeconv(kernel=4, stride=2, padding=1), torch.nn.ReLU() + torch.nn.Conv2d(8, 8, 3, padding=1), + self.DwDeconv(kernel=4, stride=2, padding=1), + torch.nn.ReLU(), ) inputs = (torch.randn(1, 8, 6, 10),) + default_passes = XNNPACKPassManager( + Tester(module.eval(), inputs).export().get_artifact() + ).passes + self.assertLess( + default_passes.index(DecomposeDepthwiseConvTransposePass), + default_passes.index(FuseActivationPass), + ) + tester = ( + Tester(module.eval(), inputs) + .export() + .to_edge() + .run_passes(RunPasses(default_passes)) + .check_count({self.conv: 5, self.pad: 4, self.d2s: 1}) + ) + graph = ( + tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph + ) + (relu,) = [ + n for n in graph.nodes if n.target == exir_ops.edge.aten.relu.default + ] + self.assertEqual( + relu.args[0].target, exir_ops.edge.et_xnnpack.depth_to_space_2d.default + ) + convs = [ + n for n in graph.nodes if n.target == exir_ops.edge.aten.convolution.default + ] + self.assertEqual( + [FuseActivationPass.get_fused_activation(n) for n in convs], [None] * 5 + ) + self.assertEqual( + sum(FusePadIntoConvPass.get_fused_padding(n) is not None for n in convs), 4 + ) + self.assertEqual( + sum( + FusePadIntoConvPass.is_folded_pad(n) + for n in graph.nodes + if n.target == exir_ops.edge.aten.constant_pad_nd.default + ), + 4, + ) self._check_lowered(module, inputs) + class ExportWithExternalConstants(BaseStages.Export): + def run(self, artifact, inputs): + tagged = torch.export.export(artifact, inputs, strict=True).module() + delegate_external_constants_pass_unlifted( + module=tagged, gen_tag_fn=lambda _: "model" + ) + self.exported_program = torch.export.export(tagged, inputs, strict=True) + + def test_external_constants(self): + # The phase weights inherit the original weight's external-constant + # tag, so every constant ends up in model.ptd and none in the .pte. + module = self.DwDeconv(kernel=4, stride=2, padding=1) + inputs = (torch.randn(1, 8, 6, 10),) + tester = Tester(module.eval(), inputs) + tester.export(self.ExportWithExternalConstants()) + tester.to_edge_transform_and_lower().check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ).to_executorch() + exec = tester.get_artifact() + self.assertIsNone(_deserialize_pte_binary(exec.buffer).named_data) + self.assertEqual(list(exec._tensor_data), ["model"]) + ptd = _deserialize_to_flat_tensor(bytes(exec._tensor_data["model"])) + self.assertEqual(len(ptd.named_data), 5) # 4 phase weights + bias + + from executorch.extension.pybindings import portable_lib as runtime + + loaded = runtime._load_for_executorch_from_buffer( + exec.buffer, bytes(exec._tensor_data["model"]) + ) + self.assertTrue( + torch.allclose(loaded.forward(inputs)[0], module(*inputs), atol=1e-3) + ) + def test_dynamic_shape(self): module = self.DwDeconv(kernel=4, stride=2, padding=1) inputs = (torch.randn(1, 8, 6, 10),) From d7d0a0b7b21fb42809c9bb94b87b7b96aaa5bcd9 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 15:57:53 -0800 Subject: [PATCH 10/15] Add Phonon (Gradium 31M TTS) example with multi-backend export (#24810) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24810 Adds `examples/models/phonon`: Gradium's Apache-2.0 Phonon 31M English text-to-speech model exported to one `.pte`, with a backend-agnostic C++ runner (`phonon_runner`) and a Python runner (`run_phonon.py`). It follows `nemotron3_diarization`: one export script, CMake presets with `make phonon-{cpu,cuda,mlx,vulkan}`, CI jobs in `pull.yml`, `mlx.yml` and `cuda.yml`, and a short README in the same shape. CUDA writes `aoti_cuda_blob.ptd` beside the `.pte`. Validated against the official runtime on XNNPACK, Vulkan, WebGPU, CUDA and MLX. Exports default to FP32 on every backend; `--dtype bf16` is available on CUDA and MLX. WebGPU runs in the browser: the README gives the `--backend webgpu --webgpu-compat-ops` export and links the Text to Speech demo of the ExecuTorch WebGPU Playground. Also fixes an XNNPACK delegate use-after-free that Phonon hits: `XNNCompiler::compileModel` frees the named-data buffer of every packed filter/bias right after `xnn_create_runtime_v4`, even when another node still reads that constant. Phonon's `StreamingConvTranspose1d` feeds the transposed-conv bias to the conv and to the held-back tail subtraction, so a program loaded from a file (C++ runner, pybindings) segfaulted in the first `mimi_decode` call. A packed value that another node also reads now stays alive; an unknown node type keeps every constant alive. Same class of bug as the PReLU fix (#21480). Differential Revision: D128116963 --- .github/workflows/cuda.yml | 48 + .github/workflows/mlx.yml | 36 + .github/workflows/pull.yml | 41 + Makefile | 26 + README.md | 2 +- backends/xnnpack/runtime/XNNCompiler.cpp | 135 +++ backends/xnnpack/test/ops/test_conv1d.py | 47 + examples/models/phonon/CMakeLists.txt | 114 ++ examples/models/phonon/CMakePresets.json | 142 +++ examples/models/phonon/NOTICE | 44 + examples/models/phonon/README.md | 91 ++ examples/models/phonon/__init__.py | 5 + examples/models/phonon/export_phonon.py | 366 ++++++ examples/models/phonon/main.cpp | 158 +++ examples/models/phonon/model.py | 1015 +++++++++++++++++ examples/models/phonon/run_phonon.py | 307 +++++ examples/models/phonon/runner.cpp | 511 +++++++++ examples/models/phonon/runner.h | 111 ++ .../models/phonon/source_transformations.py | 231 ++++ examples/models/phonon/tests/__init__.py | 5 + examples/models/phonon/tests/reference.py | 162 +++ examples/models/phonon/tests/test_phonon.py | 409 +++++++ 22 files changed, 4005 insertions(+), 1 deletion(-) create mode 100644 examples/models/phonon/CMakeLists.txt create mode 100644 examples/models/phonon/CMakePresets.json create mode 100644 examples/models/phonon/NOTICE create mode 100644 examples/models/phonon/README.md create mode 100644 examples/models/phonon/__init__.py create mode 100644 examples/models/phonon/export_phonon.py create mode 100644 examples/models/phonon/main.cpp create mode 100644 examples/models/phonon/model.py create mode 100644 examples/models/phonon/run_phonon.py create mode 100644 examples/models/phonon/runner.cpp create mode 100644 examples/models/phonon/runner.h create mode 100644 examples/models/phonon/source_transformations.py create mode 100644 examples/models/phonon/tests/__init__.py create mode 100644 examples/models/phonon/tests/reference.py create mode 100644 examples/models/phonon/tests/test_phonon.py diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 56bf2f71262..4675e757c6f 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -41,6 +41,7 @@ on: - .ci/scripts/tests/test_muse_glimmer_batching_workload.py - .ci/scripts/tests/muse_glimmer_batching_workload_test.cpp - examples/models/muse-glimmer/** + - examples/models/phonon/** - extension/pybindings/** - runtime/__init__.py - runtime/cuda.py @@ -807,6 +808,53 @@ jobs: ctest --test-dir cmake-out -R test_cuda_guard --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_stream_guard --output-on-failure -V + test-cuda-phonon: + name: test-cuda-phonon + needs: [changed-files, run-decision] + if: | + contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || + contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || + contains(needs.changed-files.outputs.changed-files, 'examples/models/phonon') || + contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda.yml') || + needs.run-decision.outputs.is-full-run == 'true' + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read + with: + timeout: 90 + runner: mt-l-x86aavx2-11-41-a10g + gpu-arch-type: cuda + gpu-arch-version: "13.0" + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + script: | + set -eux + + # OSDC mounts HF_HOME read-only at /mnt/hf_cache; redirect to a writable dir + # (RUNNER_TEMP, or /tmp when RUNNER_TEMP isn't writable inside the container). + export HF_HOME="${RUNNER_TEMP:-/tmp}/hf_cache" + mkdir -p "${HF_HOME}" 2>/dev/null || export HF_HOME=/tmp/hf_cache + mkdir -p "${HF_HOME}" + # The runner also points HF_HUB_CACHE at the read-only mount. + export HF_HUB_CACHE="${HF_HOME}/hub" + + bash ./install_executorch.sh + conda install -y -c conda-forge 'libstdcxx-ng>=12' + export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + pip install safetensors tokenizers huggingface_hub + python -m pytest examples/models/phonon/tests -q + python -m executorch.examples.models.phonon.export_phonon \ + --backend cuda --output-dir phonon_exports/cuda + make phonon-cuda + cmake-out/examples/models/phonon/phonon_runner \ + --model_path phonon_exports/cuda/phonon.pte \ + --data_path phonon_exports/cuda/aoti_cuda_blob.ptd \ + --text "Hello from ExecuTorch." \ + --output phonon.wav + # 16-bit mono WAV with at least 0.5 s of audio at 24 kHz. + test "$(wc -c < phonon.wav)" -gt $((44 + 24000)) + test-model-cuda-e2e: name: test-model-cuda-e2e-${{ matrix.model.name }}-${{ matrix.quant }} # Skip this job if the pull request is from a fork (HuggingFace secrets are not available). diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index fa586b84831..b18758f22fb 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -24,6 +24,7 @@ on: - examples/models/muse-glimmer/** - examples/models/nemotron3_diarization/** - examples/models/parakeet/** + - examples/models/phonon/** - examples/models/supertonic/** - examples/models/voxtral_realtime/** - examples/models/qwen3_5_moe/** @@ -409,6 +410,41 @@ jobs: test -f cmake-out/lib/libextension_asr_runner.a test -f cmake-out/lib/libextension_llm_runner.a + test-mlx-phonon: + needs: run-decision + if: | + github.event_name == 'pull_request' || + needs.run-decision.outputs.is-full-run == 'true' + uses: pytorch/test-infra/.github/workflows/macos_job.yml@main + with: + default-packages: "" + job-name: test-mlx-phonon + runner: macos-15-xlarge + python-version: "3.12" + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -eux + # Top-level presets build with unlimited make -j and per-model presets serially. + export CMAKE_BUILD_PARALLEL_LEVEL=$(( $(sysctl -n hw.ncpu) - 1 )) + # Disable HF Xet storage to avoid stalled downloads on CI runners + export HF_HUB_DISABLE_XET=1 + ${CONDA_RUN} python install_executorch.py > /dev/null + ${CONDA_RUN} pip install safetensors tokenizers huggingface_hub + ${CONDA_RUN} python -m pytest examples/models/phonon/tests -q + ${CONDA_RUN} python -m executorch.examples.models.phonon.export_phonon \ + --backend mlx --output-dir phonon_exports/mlx + ${CONDA_RUN} make phonon-mlx + BIN_DIR=cmake-out/examples/models/phonon + test -f "${BIN_DIR}/mlx.metallib" + "${BIN_DIR}/phonon_runner" \ + --model_path phonon_exports/mlx/phonon.pte \ + --text "Hello from ExecuTorch." \ + --output phonon.wav + # 16-bit mono WAV with at least 0.5 s of audio at 24 kHz. + test "$(wc -c < phonon.wav)" -gt $((44 + 24000)) + backend-tester: needs: run-decision if: | diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 2ea96fe18e0..daaf1ebb0a5 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -366,6 +366,47 @@ jobs: bash .ci/scripts/test_model_e2e.sh xnnpack "nvidia/parakeet-tdt" "quantized-8da4w" "./parakeet_output" echo "::endgroup::" + test-phonon-xnnpack-linux: + needs: docker-image + name: test-phonon-xnnpack-linux + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read + with: + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -eux + + # The generic Linux job chooses to use base env, not the one setup by the image + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + + # OSDC mounts HF_HOME read-only at /mnt/hf_cache; redirect to a writable dir + # (RUNNER_TEMP, or /tmp when RUNNER_TEMP isn't writable inside the container). + export HF_HOME="${RUNNER_TEMP:-/tmp}/hf_cache" + mkdir -p "${HF_HOME}" 2>/dev/null || export HF_HOME=/tmp/hf_cache + mkdir -p "${HF_HOME}" + # The runner also points HF_HUB_CACHE at the read-only mount. + export HF_HUB_CACHE="${HF_HOME}/hub" + + ./install_executorch.sh + pip install safetensors tokenizers huggingface_hub + python -m pytest examples/models/phonon/tests -q + python -m executorch.examples.models.phonon.export_phonon \ + --backend xnnpack --output-dir phonon_exports/xnnpack + make phonon-cpu + cmake-out/examples/models/phonon/phonon_runner \ + --model_path phonon_exports/xnnpack/phonon.pte \ + --text "Hello from ExecuTorch." \ + --output phonon.wav + # 16-bit mono WAV with at least 0.5 s of audio at 24 kHz. + test "$(wc -c < phonon.wav)" -gt $((44 + 24000)) + test-voxtral-realtime-xnnpack-linux: needs: docker-image name: test-voxtral-realtime-xnnpack-linux diff --git a/Makefile b/Makefile index a021214ab7b..0ac6073496d 100644 --- a/Makefile +++ b/Makefile @@ -20,6 +20,7 @@ # - parakeet: Speech recognition model (CPU, CUDA, Metal, MLX) # - sortformer: Speaker diarization model (CPU, CUDA) # - nemotron3-diarization: Speaker diarization model (MLX, XNNPACK, CUDA, Vulkan) +# - phonon: Text-to-speech model (XNNPACK, CUDA, MLX, Vulkan) # - supertonic: Text-to-speech model (MLX) # - silero_vad: Voice activity detection model (CPU) # - llama: Text generation model (CPU) @@ -128,6 +129,10 @@ help: @echo " nemotron3-diarization-cpu - Build Nemotron 3 Diarization runner with XNNPACK backend" @echo " nemotron3-diarization-cuda - Build Nemotron 3 Diarization runner with CUDA backend" @echo " nemotron3-diarization-vulkan - Build Nemotron 3 Diarization runner with Vulkan backend" + @echo " phonon-cpu - Build Phonon runner with XNNPACK backend" + @echo " phonon-cuda - Build Phonon runner with CUDA backend" + @echo " phonon-mlx - Build Phonon runner with MLX backend" + @echo " phonon-vulkan - Build Phonon runner with Vulkan backend" @echo " supertonic-mlx - Build Supertonic runner with MLX backend" @echo " silero-vad-cpu - Build Silero VAD runner with CPU backend" @echo " llama-cuda - Build Llama runner with CUDA backend" @@ -317,6 +322,27 @@ nemotron3-diarization-vulkan: cmake --build --preset llm-release-install cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-vulkan +.PHONY: phonon-cpu +phonon-cpu: + cmake --workflow --preset llm-release + cd examples/models/phonon && cmake --workflow --preset phonon-cpu + +.PHONY: phonon-cuda +phonon-cuda: + cmake --workflow --preset llm-release-cuda + cd examples/models/phonon && cmake --workflow --preset phonon-cuda + +.PHONY: phonon-mlx +phonon-mlx: + cmake --workflow --preset mlx-release + cd examples/models/phonon && cmake --workflow --preset phonon-mlx + +.PHONY: phonon-vulkan +phonon-vulkan: + cmake --preset llm-release -DEXECUTORCH_BUILD_VULKAN=ON + cmake --build --preset llm-release-install + cd examples/models/phonon && cmake --workflow --preset phonon-vulkan + sortformer-cuda: @echo "==> Building and installing ExecuTorch with CUDA..." cmake --workflow --preset llm-release-cuda diff --git a/README.md b/README.md index b5ce3d55582..ac67afe8660 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ and Messenger**, serving billions of people. It also runs AI features on | Workload | Representative capabilities | Start here | |---|---|---| | **Local LLM and agent building blocks** | Quantized text models, long context, tool calling, speculative decoding, multi-session execution, and an experimental OpenAI-compatible local server | [Muse Glimmer](examples/models/muse-glimmer/README.md) · [LLM guide](https://docs.pytorch.org/executorch/main/llm/working-with-llms.html) · [LLM server](examples/llm_server/README.md) | -| **Voice** | Streaming and offline speech recognition, speech synthesis, voice activity detection, and speaker diarization | [Voxtral Realtime](examples/models/voxtral_realtime/README.md) · [Parakeet](examples/models/parakeet/README.md) · [Sortformer diarization](examples/models/sortformer/README.md) · [Voxtral TTS](examples/models/voxtral_tts/README.md) · [Nemotron 3 Diarization](examples/models/nemotron3_diarization/README.md) | +| **Voice** | Streaming and offline speech recognition, speech synthesis, voice activity detection, and speaker diarization | [Voxtral Realtime](examples/models/voxtral_realtime/README.md) · [Parakeet](examples/models/parakeet/README.md) · [Sortformer diarization](examples/models/sortformer/README.md) · [Voxtral TTS](examples/models/voxtral_tts/README.md) · [Nemotron 3 Diarization](examples/models/nemotron3_diarization/README.md) · [Phonon TTS](examples/models/phonon/README.md) | | **Multimodal** | Text, image, and audio runners; mobile vision-language models | [Gemma 4](examples/models/gemma4/README.md) · [Multimodal runner](extension/llm/runner/README.md) | | **Computer vision** | Image classification, object detection, semantic segmentation, and promptable segmentation | [MobileNet V2](https://docs.pytorch.org/executorch/main/getting-started.html#preparing-the-model) · [YOLO26](examples/models/yolo26/README.md) · [DeepLabV3](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo) · [EfficientSAM](examples/models/efficient_sam/README.md) | | **Embedded AI** | Cortex-M CPU kernels, Ethos-U and NXP NPUs, Cadence DSPs, Zephyr, Arduino, and Raspberry Pi Pico workflows | [Embedded guide](https://docs.pytorch.org/executorch/main/embedded-section.html) · [Cortex-M](https://docs.pytorch.org/executorch/main/backends/arm-cortex-m/arm-cortex-m-overview.html) · [Arduino](examples/arduino/README.md) | diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index b0efbd0975d..380d154ea22 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -2283,6 +2284,126 @@ uint32_t getSerializedValueId(ValuePtr value) noexcept { : XNN_INVALID_VALUE_ID; } +/* +Appends the ids of the values `node` reads, except the filter and bias of +FullyConnected and the convolutions, which XNNPACK packs. Returns false for a +node type it does not know. +*/ +bool appendUnpackedInputIds( + const fb_xnnpack::XNode* node, + std::vector& ids) { + using U = fb_xnnpack::XNodeUnion; + const void* p = node->xnode_union(); + if (p == nullptr) { + return false; + } + const U type = node->xnode_union_type(); + auto any_of = [type](std::initializer_list types) { + for (U t : types) { + if (t == type) { + return true; + } + } + return false; + }; + if (any_of( + {U::XNNAdd, + U::XNNDiv, + U::XNNMinimum, + U::XNNMultiply, + U::XNNSubtract, + U::XNNMaximum, + U::XNNPReLU, + U::XNNBatchMatrixMultiply})) { + auto n = static_cast(p); + ids.push_back(n->input1_id()); + ids.push_back(n->input2_id()); + } else if (any_of( + {U::XNNSoftmax, + U::XNNSigmoid, + U::XNNClamp, + U::XNNFloor, + U::XNNConvert, + U::XNNGlobalAvgPooling2d, + U::XNNSquareRoot, + U::XNNCeiling, + U::XNNHardswish, + U::XNNNegate, + U::XNNSquare, + U::XNNAbs, + U::XNNReciprocalSquareRoot, + U::XNNLog, + U::XNNGelu, + U::XNNTanh, + U::XNNExp, + U::XNNSin, + U::XNNCopy, + U::XNNCos, + U::XNNApproxGelu})) { + ids.push_back(static_cast(p)->input_id()); + } else if (any_of( + {U::XNNConcatenate2, + U::XNNConcatenate3, + U::XNNConcatenate4, + U::XNNConcatenate5})) { + auto n = static_cast(p); + ids.insert( + ids.end(), + {n->input1_id(), + n->input2_id(), + n->input3_id(), + n->input4_id(), + n->input5_id()}); + } else if (type == U::XNNFullyConnected) { + ids.push_back( + static_cast(p)->input1_id()); + } else if (any_of( + {U::XNNConv2d, + U::XNNDepthwiseConv2d, + U::XNNConvTranspose2d})) { + ids.push_back(static_cast(p)->input1_id()); + } else if (any_of({U::XNNAvgPooling2d, U::XNNMaxPooling2d})) { + ids.push_back(static_cast(p)->input_id()); + } else if (type == U::XNNStaticTranspose) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNStaticResizeBilinear2D) { + ids.push_back(static_cast(p) + ->input_id()); + } else if (type == U::XNNStaticConstantPad) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNStaticReshape) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNStaticSlice) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNDepthToSpace2d) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNELU) { + ids.push_back(static_cast(p)->input_id()); + } else if (type == U::XNNLeakyReLU) { + ids.push_back(static_cast(p)->input_id()); + } else if (type == U::XNNArgMaxPooling2d) { + ids.push_back( + static_cast(p)->input_id()); + } else if (type == U::XNNScaledDotProductAttention) { + auto n = static_cast(p); + ids.insert( + ids.end(), + {n->query_id(), + n->key_id(), + n->value_id(), + n->scale_id(), + n->mask_id()}); + } else { + return false; + } + return true; +} + /* Serialized ids of the values that XNNPACK copies into its own packed storage while the runtime is created. Every other constant value is referenced by raw @@ -2315,6 +2436,20 @@ std::unordered_set getPackedValueIds(GraphPtr flatbuffer_graph) { insert_weights(n->filter_id(), n->bias_id()); } } + + // A packed weight or bias that another node also reads (a bias reused by a + // reshape or a subtract, for example) is referenced by raw pointer by that + // node, so it must stay alive with the runtime. + std::vector reads; + for (auto node : *flatbuffer_graph->xnodes()) { + if (!appendUnpackedInputIds(node, reads)) { + // A node type listed nowhere below: keep every constant alive. + return {}; + } + } + for (uint32_t id : reads) { + packed_ids.erase(id); + } return packed_ids; } diff --git a/backends/xnnpack/test/ops/test_conv1d.py b/backends/xnnpack/test/ops/test_conv1d.py index 3ff32354167..e03341870dc 100644 --- a/backends/xnnpack/test/ops/test_conv1d.py +++ b/backends/xnnpack/test/ops/test_conv1d.py @@ -4,9 +4,13 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import gc +import tempfile import unittest +from pathlib import Path import torch +from executorch.backends.test.harness.stages import StageType from executorch.backends.xnnpack.partition.config.xnnpack_config import ( ConfigPrecisionType, ) @@ -16,6 +20,7 @@ from executorch.backends.xnnpack.test.tester import Quantize, RunPasses, Tester from executorch.backends.xnnpack.test.tester.tester import ToEdgeTransformAndLower from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from executorch.runtime import Runtime, Verification class TestConv1d(unittest.TestCase): @@ -292,3 +297,45 @@ def test_qs8_conv1d_with_floating_point_partitioner(self): passes=RunPasses(pass_functions=[constant_prop_pass]), skip_to_executorch=True, ) + + class ConvTranspose1dBiasReused(torch.nn.Module): + def __init__(self): + super().__init__() + self.convtr = torch.nn.ConvTranspose1d(8, 4, kernel_size=4, stride=2) + + def forward(self, x): + return self.convtr(x) - self.convtr.bias.view(1, -1, 1) + + def test_fp32_conv_transpose1d_bias_reused_file_load(self): + """ + XNNPACK packs a convolution's bias, so the delegate frees its data after + init. Here the subtraction reads the same bias by pointer at every + execution, so it must stay alive. Only a load from a file really frees + the data. + """ + module = self.ConvTranspose1dBiasReused().eval() + x = torch.randn(1, 8, 6) + expected = module(x) + + tester = ( + Tester(module, (x,)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + ) + buf = tester.stages[StageType.SERIALIZE].artifact + with tempfile.TemporaryDirectory() as temp_dir: + pte_path = Path(temp_dir) / "conv_transpose1d.pte" + pte_path.write_bytes(buf) + program = Runtime.get().load_program( + pte_path, verification=Verification.Minimal + ) + method = program.load_method("forward") + actual = method.execute((x,))[0].clone() + # Release the mapping before the temp dir is removed (Windows). + del method, program + gc.collect() + + self.assertTrue(torch.allclose(expected, actual, atol=1e-5)) diff --git a/examples/models/phonon/CMakeLists.txt b/examples/models/phonon/CMakeLists.txt new file mode 100644 index 00000000000..c6f32c504f8 --- /dev/null +++ b/examples/models/phonon/CMakeLists.txt @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +cmake_minimum_required(VERSION 3.24) +project(phonon_runner) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) +include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) + +set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags) +find_package(gflags REQUIRED) +find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH) +if(NOT TARGET tokenizers::tokenizers) + message( + FATAL_ERROR + "Build and install ExecuTorch with EXECUTORCH_BUILD_EXTENSION_LLM=ON" + ) +endif() + +add_library(phonon STATIC runner.cpp) +target_include_directories( + phonon + PUBLIC ${EXECUTORCH_ROOT}/.. + PRIVATE ${EXECUTORCH_ROOT}/third-party/json/include +) +target_link_libraries( + phonon + PUBLIC executorch + extension_module + extension_data_loader + extension_tensor + extension_flat_tensor + optimized_native_cpu_ops_lib + cpublas + eigen_blas + tokenizers::tokenizers +) +executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib) + +# Kernels for the ops that stay on CPU: the flow LM's KV cache attention +# (llama::sdpa_with_kv_cache) and --quantize quantized ops. +foreach(lib custom_ops quantized_ops_lib) + if(TARGET ${lib}) + target_link_libraries(phonon PUBLIC ${lib}) + executorch_target_link_options_shared_lib(${lib}) + endif() +endforeach() + +if(TARGET xnnpack_backend) + target_link_libraries( + phonon PUBLIC xnnpack_backend XNNPACK xnnpack-microkernels-prod + ) + if(TARGET kleidiai) + target_link_libraries(phonon PUBLIC kleidiai) + endif() + executorch_target_link_options_shared_lib(xnnpack_backend) +endif() +if(EXECUTORCH_BUILD_CUDA) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(phonon PUBLIC aoti_cuda_backend) + if(NOT MSVC) + executorch_target_link_options_shared_lib(aoti_cuda_backend) + endif() +endif() +if(EXECUTORCH_BUILD_VULKAN) + if(NOT TARGET vulkan_backend) + message( + FATAL_ERROR "Build and install ExecuTorch with EXECUTORCH_BUILD_VULKAN=ON" + ) + endif() + executorch_target_link_options_shared_lib(vulkan_backend) + target_link_libraries(phonon PUBLIC vulkan_backend) +endif() + +if(EXECUTORCH_BUILD_WEBGPU) + # The installed webgpu_backend links dawn::webgpu_dawn; pass -DDawn_DIR. + find_package(Dawn REQUIRED) + if(NOT TARGET webgpu_backend) + message( + FATAL_ERROR "Build and install ExecuTorch with EXECUTORCH_BUILD_WEBGPU=ON" + ) + endif() + executorch_target_link_options_shared_lib(webgpu_backend) + target_link_libraries(phonon PUBLIC webgpu_backend) +endif() + +add_executable(phonon_runner main.cpp) +target_link_libraries(phonon_runner PRIVATE phonon gflags) +if(EXECUTORCH_BUILD_CUDA) + target_compile_definitions(phonon_runner PRIVATE EXECUTORCH_BUILD_CUDA) +endif() +if(EXECUTORCH_BUILD_MLX) + if(NOT TARGET mlxdelegate OR NOT TARGET mlx) + message( + FATAL_ERROR "Build and install ExecuTorch with EXECUTORCH_BUILD_MLX=ON" + ) + endif() + executorch_target_link_options_shared_lib(mlxdelegate) + target_link_libraries(phonon_runner PRIVATE mlxdelegate mlx) + executorch_target_copy_mlx_metallib(phonon_runner) +endif() +if(MSVC AND EXECUTORCH_BUILD_CUDA) + add_custom_command( + TARGET phonon_runner + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + $ $ + ) +endif() diff --git a/examples/models/phonon/CMakePresets.json b/examples/models/phonon/CMakePresets.json new file mode 100644 index 00000000000..afc233e6253 --- /dev/null +++ b/examples/models/phonon/CMakePresets.json @@ -0,0 +1,142 @@ +{ + "version": 6, + "configurePresets": [ + { + "name": "phonon-base", + "hidden": true, + "binaryDir": "${sourceDir}/../../../cmake-out/examples/models/phonon", + "cacheVariables": { + "EXECUTORCH_BUILD_CUDA": "OFF", + "EXECUTORCH_BUILD_MLX": "OFF", + "CMAKE_BUILD_TYPE": "Release", + "CMAKE_FIND_ROOT_PATH": "${sourceDir}/../../../cmake-out", + "CMAKE_PREFIX_PATH": "${sourceDir}/../../../cmake-out", + "EXECUTORCH_BUILD_VULKAN": "OFF" + } + }, + { + "name": "phonon-mlx", + "displayName": "Phonon runner (MLX)", + "inherits": [ + "phonon-base" + ], + "cacheVariables": { + "EXECUTORCH_BUILD_MLX": "ON", + "CMAKE_OSX_DEPLOYMENT_TARGET": "14.0" + }, + "condition": { + "type": "equals", + "lhs": "${hostSystemName}", + "rhs": "Darwin" + } + }, + { + "name": "phonon-cuda", + "displayName": "Phonon runner (CUDA)", + "inherits": [ + "phonon-base" + ], + "cacheVariables": { + "EXECUTORCH_BUILD_CUDA": "ON" + }, + "condition": { + "type": "inList", + "string": "${hostSystemName}", + "list": ["Linux", "Windows"] + } + }, + { + "name": "phonon-cpu", + "displayName": "Phonon runner (XNNPACK)", + "inherits": [ + "phonon-base" + ] + }, + { + "name": "phonon-vulkan", + "displayName": "Phonon runner (Vulkan)", + "inherits": [ + "phonon-base" + ], + "cacheVariables": { + "EXECUTORCH_BUILD_VULKAN": "ON" + } + } + ], + "buildPresets": [ + { + "name": "phonon-cuda", + "configurePreset": "phonon-cuda", + "configuration": "Release" + }, + { + "name": "phonon-mlx", + "configurePreset": "phonon-mlx", + "configuration": "Release" + }, + { + "name": "phonon-cpu", + "configurePreset": "phonon-cpu", + "configuration": "Release" + }, + { + "name": "phonon-vulkan", + "configurePreset": "phonon-vulkan", + "configuration": "Release" + } + ], + "workflowPresets": [ + { + "name": "phonon-cuda", + "steps": [ + { + "type": "configure", + "name": "phonon-cuda" + }, + { + "type": "build", + "name": "phonon-cuda" + } + ] + }, + { + "name": "phonon-mlx", + "steps": [ + { + "type": "configure", + "name": "phonon-mlx" + }, + { + "type": "build", + "name": "phonon-mlx" + } + ] + }, + { + "name": "phonon-cpu", + "steps": [ + { + "type": "configure", + "name": "phonon-cpu" + }, + { + "type": "build", + "name": "phonon-cpu" + } + ] + }, + { + "name": "phonon-vulkan", + "steps": [ + { + "type": "configure", + "name": "phonon-vulkan" + }, + { + "type": "build", + "name": "phonon-vulkan" + } + ] + } + ] +} diff --git a/examples/models/phonon/NOTICE b/examples/models/phonon/NOTICE new file mode 100644 index 00000000000..4ef853e2dae --- /dev/null +++ b/examples/models/phonon/NOTICE @@ -0,0 +1,44 @@ +Phonon upstream attribution +=========================== + +The Phonon model, its weights, tokenizer and voices are by Gradium: + + https://huggingface.co/gradium/phonon-31m-en + Revision: 52433c46975bfa13144c56204c9f4ee1c5b148aa + License: Apache-2.0 + +Phonon builds on Pocket TTS and the Mimi codec, developed by Kyutai. + +The PyTorch model in model.py and the generation policy in run_phonon.py +(text preparation, frame budget and end-of-speech rule) follow the official +Phonon runtime, xn-ptts: + + https://github.com/gradium-ai/xn-ptts + Version: ptts 0.4.0 (crates.io) + +xn-ptts is licensed MIT OR Apache-2.0. This example uses it under the MIT +license: + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md new file mode 100644 index 00000000000..2c7f568ed3f --- /dev/null +++ b/examples/models/phonon/README.md @@ -0,0 +1,91 @@ +# Phonon + +Text-to-speech at 24 kHz with Gradium's +[Phonon 31M English](https://huggingface.co/gradium/phonon-31m-en), built on +Kyutai's [Pocket TTS](https://github.com/kyutai-labs/pocket-tts) and the Mimi +codec. The weights, tokenizer and voices are Apache-2.0. + +```bash +# From the repository root, in a new Python 3.11+ environment. Install +# ExecuTorch from this checkout; it also builds CUDA if nvcc is on PATH, +# Vulkan if glslc is on PATH, and MLX if `xcrun --find metal` succeeds. +./install_executorch.sh +python -m pip install safetensors tokenizers huggingface_hub + +# Choose xnnpack (CPU), cuda (NVIDIA GPU), mlx (Apple Silicon), or vulkan (GPU). +# For webgpu (the browser), see WebGPU below. +BACKEND=xnnpack +# Export: FP32 for every backend. Add --dtype bf16 for BF16 on CUDA or MLX. +python -m executorch.examples.models.phonon.export_phonon \ + --backend "$BACKEND" \ + --output-dir "phonon_exports/$BACKEND" + +# Build the matching runner (Vulkan requires the Vulkan SDK and glslc). +# Remove cmake-out before building for another backend. +make "phonon-${BACKEND/xnnpack/cpu}" + +# Synthesize +DATA_PATH="" +if [ "$BACKEND" = cuda ]; then + DATA_PATH="phonon_exports/cuda/aoti_cuda_blob.ptd" +fi +cmake-out/examples/models/phonon/phonon_runner \ + --model_path="phonon_exports/$BACKEND/phonon.pte" \ + --data_path="$DATA_PATH" \ + --text="Hello from ExecuTorch." \ + --output=hello.wav +# Prints "Wrote hello.wav: 2.16 s of audio, 27 frames, ..." for an FP32 export. +``` + +Export downloads `gradium/phonon-31m-en` at a pinned revision (or reads +`--checkpoint-dir`) and writes `phonon.pte`, `phonon.json` (the inputs and +outputs of its four methods), the tokenizer and the default voice. +`--voice_path` takes another voice from +[gradium/phonon-voices](https://huggingface.co/gradium/phonon-voices); +`--temperature` (0.3 by default, 0 for no noise) and `--seed` control sampling. +`run_phonon.py` is the same runner in Python and generates the same audio from +the same program (`python -m executorch.examples.models.phonon.run_phonon +--model phonon_exports/$BACKEND/phonon.pte --text "Hello from ExecuTorch." +--output hello_py.wav`, plus `--data-path` for CUDA). + +XNNPACK delegates the linear and convolution ops; the flow LM's attention runs +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. +`--quantize 4bit` (8-bit dynamic activations, 4-bit weights) or `--quantize 8bit` +quantizes the flow LM's linear layers on XNNPACK. A backend without a kernel for +the level fails before the download and names the options that work. + +Vulkan and MLX run each method in one partition. Vulkan needs `glslc`, for +example from the [Vulkan SDK](https://vulkan.lunarg.com/sdk/home) (on Linux, +extract the tarball and `source setup-env.sh`). MLX needs the full Xcode with +the Metal toolchain, selected with `xcode-select` before +`./install_executorch.sh` (see the [MLX backend](../../../backends/mlx/README.md)); +the first run compiles the Metal kernels. + +CUDA export requires a CUDA toolkit; with its `nvcc` on `PATH`, +`./install_executorch.sh` installs CUDA-enabled PyTorch and the CUDA backend. +Each method runs in one CUDA partition, with float32 inputs and outputs and +FP32 (or BF16) model computation. Export produces `phonon.pte` and +`aoti_cuda_blob.ptd`; pass the latter through `--data_path` as shown above. If +the runner cannot load a method on a system with an older libstdc++, put the +exporting Python environment's libraries first: +`export LD_LIBRARY_PATH="$(python -c 'import sys; print(sys.prefix)')/lib:$LD_LIBRARY_PATH"`. + +BF16 halves the weights, but small differences grow through the autoregressive +loop into a different take of the same text; FP32 matches the official runtime +sample for sample. + +WebGPU runs Phonon in the browser on the ExecuTorch Wasm runtime; try it in +the Text to Speech demo of the +[ExecuTorch WebGPU Playground](https://executorch.ai/playground/). Export with +`--backend webgpu --webgpu-compat-ops --output-dir phonon_exports/webgpu`. +The WebGPU backend does not delegate `aten.elu` or 1-D transposed convolution +yet; `--webgpu-compat-ops` replaces them (and `native_layer_norm` and +`aten.silu`) with equivalent ops, so every method but the text embedding runs +on the GPU. Without it, those ops run on the CPU. + +Phonon reads English only; spell out numbers and symbols. The tests +(`python -m pytest examples/models/phonon/tests -q`) compare the exported +programs of a small random model with a reference implementation of the +official runtime (see [NOTICE](NOTICE)). To check your own export, synthesize +"The quick brown fox jumps over the lazy dog." with `--temperature=0`: every +FP32 backend gives 37 frames (2.96 s), as the official runtime does. diff --git a/examples/models/phonon/__init__.py b/examples/models/phonon/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/examples/models/phonon/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py new file mode 100644 index 00000000000..0bf5f5181fa --- /dev/null +++ b/examples/models/phonon/export_phonon.py @@ -0,0 +1,366 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Export Gradium's Phonon text-to-speech model to one ExecuTorch program. + +The .pte holds four methods (text_embed, voice_embed, flow_lm, mimi_decode) +and constant metadata methods. The I/O contract is the same for every backend +and is also written to phonon.json next to the program. + + python -m executorch.examples.models.phonon.export_phonon \\ + --backend xnnpack --output-dir phonon_exports/xnnpack +""" + +import argparse +import json +import shutil +from pathlib import Path +from typing import Any + +import torch +from executorch.examples.models.phonon.model import ( + build_programs, + dynamic_shapes, + example_inputs, + io_contract, + mimi_state_spec, + Phonon, + program_metadata, +) +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.exir.passes import MemoryPlanningPass +from torch.export import export + +HF_MODEL = "gradium/phonon-31m-en" +HF_REVISION = "52433c46975bfa13144c56204c9f4ee1c5b148aa" +CHECKPOINT_FILES = ( + "model.safetensors", + "config.json", + "tokenizer.json", + "tokenizer.model", + "default-voice.safetensors", +) +BACKENDS = ("portable", "xnnpack", "webgpu", "vulkan", "cuda", "mlx") +# Methods with a dynamic sequence length. +DYNAMIC_METHODS = ("text_embed", "voice_embed", "flow_lm") +DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16} +# Backends that support --dtype bf16. +BF16_BACKENDS = ("portable", "cuda", "mlx") +# The linear quantization scheme of each --quantize level, per backend: 8da is +# dynamically quantized 8-bit activations with 4- or 8-bit weights. +QUANTIZE_SCHEMES = {"xnnpack": {"4bit": "8da4w", "8bit": "8da8w"}} + + +def attention_for(backend: str) -> str: + """KV cache / attention implementation of the flow LM per backend (see + FlowLMStep): the ExecuTorch custom op where a kernel exists for it, plain + SDPA with an explicit mask for the compiler backends.""" + return "sdpa" if backend in ("cuda", "mlx") else "custom_sdpa" + + +def source_transforms_for(backend: str, webgpu_compat_ops: bool): + """Module swaps applied before export (see source_transformations.py).""" + from executorch.examples.models.phonon import source_transformations as st + + if webgpu_compat_ops: + if backend != "webgpu": + raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") + return [st.use_webgpu_compat_ops] + return { + "mlx": [st.decompose_elu], + "vulkan": [ + st.decompose_elu, + st.depthwise_conv_transpose_as_mul, + st.conv_transpose1d_as_2d, + ], + }.get(backend, []) + + +def quantize_scheme(backend: str, quantize: str | None) -> str | None: + """The linear quantization scheme for `--quantize` on `backend`. Raises + ValueError with what to run instead when the backend has no kernel for the + level.""" + if quantize is None: + return None + schemes = QUANTIZE_SCHEMES.get(backend, {}) + if quantize not in schemes: + options = [f"with --quantize {q}" for q in schemes] + options.append("without --quantize (full precision)") + raise ValueError( + f"--backend {backend} has no {quantize} linear kernel. Export " + f"{' or '.join(options)} instead." + ) + return schemes[quantize] + + +def partitioners_for(backend: str, methods, quantized: bool): + """Per-method partitioners, transform passes and graph decompositions for + one backend. Returns (partitioners, transform_passes, decompositions).""" + if backend == "portable": + return {m: [] for m in methods}, None, None + if backend == "xnnpack": + from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( + XnnpackDynamicallyQuantizedPartitioner, + XnnpackPartitioner, + ) + + parts = {m: [XnnpackPartitioner()] for m in methods} + if quantized: + # --quantize quantizes the flow LM only. + parts["flow_lm"].insert(0, XnnpackDynamicallyQuantizedPartitioner()) + return parts, None, None + if backend in ("vulkan", "webgpu"): + if backend == "vulkan": + from executorch.backends.vulkan.partitioner.vulkan_partitioner import ( + VulkanPartitioner as Partitioner, + ) + else: + from executorch.backends.webgpu.partitioner import ( + WebGPUPartitioner as Partitioner, + ) + parts = {} + for m in methods: + options = {"force_fp16": False} + if m in DYNAMIC_METHODS: + options["require_dynamic_shapes"] = True + parts[m] = [Partitioner(compile_options=options)] + if backend == "webgpu": + # The WebGPU runtime does not resize an embedding output to a + # dynamic token count yet; the lookup runs on the CPU kernels. + parts["text_embed"] = [] + return parts, None, None + if backend == "cuda": + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from torch._inductor.decomposition import conv1d_to_conv2d + + parts = { + m: [CudaPartitioner([CudaBackend.generate_method_name_compile_spec(m)])] + for m in methods + } + return parts, None, {torch.ops.aten.conv1d.default: conv1d_to_conv2d} + if backend == "mlx": + from executorch.backends.mlx.partitioner import MLXPartitioner + from executorch.backends.mlx.passes import get_default_passes + + return {m: [MLXPartitioner()] for m in methods}, get_default_passes(), None + raise ValueError(f"Unknown backend: {backend}") + + +def resolve_checkpoint(args) -> Path: + if args.checkpoint_dir: + return Path(args.checkpoint_dir) + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + args.hf_model, revision=args.revision, allow_patterns=list(CHECKPOINT_FILES) + ) + ) + + +def delegation_summary(edge) -> dict[str, Any]: + from executorch.devtools.backend_debug import get_delegation_info + + summary = {} + for name in edge.methods: + info = get_delegation_info(edge.exported_program(name).graph_module) + cpu_ops = sorted( + op + for op, b in info.delegation_by_operator.items() + if b.non_delegated + and op not in ("getitem", "sym_size", "_local_scalar_dense") + ) + summary[name] = { + "delegates": info.num_delegated_subgraphs, + "delegated_nodes": info.num_delegated_nodes, + "non_delegated_ops": cpu_ops, + } + return summary + + +def lower_phonon( + model: Phonon, + backend: str = "xnnpack", + quantize: str | None = None, + quantize_group_size: int | None = None, + webgpu_compat_ops: bool = False, + dtype: str = "fp32", +): + """Export and lower the four programs of `model` for `backend`. + Returns (ExecutorchProgramManager, delegation summary, metadata).""" + from executorch.examples.models.phonon.source_transformations import ( + compute_weights_in, + ) + + cfg = model.cfg + qlinear = quantize_scheme(backend, quantize) + if dtype == "bf16" and (backend not in BF16_BACKENDS or qlinear): + raise ValueError( + f"--dtype bf16 is available for --backend {', '.join(BF16_BACKENDS)} " + "without --quantize. Export without --dtype (FP32)." + ) + # Before the module swaps, which change the model's modules in place. + metadata = program_metadata(model) + programs = build_programs(model, attention_for(backend)) + for transform in source_transforms_for(backend, webgpu_compat_ops): + print(f"{transform.__name__}: {transform(programs)}") + if dtype == "bf16": + print( + f"compute_weights_in(bf16): {compute_weights_in(programs, DTYPES[dtype])}" + ) + if qlinear: + from executorch.extension.llm.export.quantize import quantize_model_ + + # Weight quantization applies to the flow LM; the codec stays in fp32. + quantize_model_( + programs["flow_lm"], + qlinear_config=qlinear, + qlinear_group_size=quantize_group_size, + skip_incompatible_shapes=True, + ) + + inputs, shapes = example_inputs(cfg), dynamic_shapes(cfg) + partitioners, passes, decompositions = partitioners_for( + backend, list(programs), quantized=bool(qlinear) + ) + eps = {} + with torch.no_grad(): + for name, program in programs.items(): + ep = export( + program, inputs[name], dynamic_shapes=shapes[name], strict=False + ) + if decompositions: + ep = ep.run_decompositions(decompositions) + eps[name] = ep + edge = to_edge_transform_and_lower( + eps, + partitioner=partitioners, + transform_passes={m: passes for m in eps} if passes else None, + constant_methods=metadata, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + delegation = delegation_summary(edge) + program = edge.to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + return program, delegation, metadata + + +def export_phonon( + checkpoint_dir: Path, + output_dir: Path, + backend: str = "xnnpack", + quantize: str | None = None, + quantize_group_size: int | None = None, + webgpu_compat_ops: bool = False, + dtype: str = "fp32", +) -> Path: + """Export a checkpoint directory to output_dir/phonon.pte (plus + aoti_cuda_blob.ptd for CUDA), phonon.json, the tokenizer and the default + voice.""" + model = Phonon.from_checkpoint(checkpoint_dir / "model.safetensors") + cfg = model.cfg + program, delegation, metadata = lower_phonon( + model, backend, quantize, quantize_group_size, webgpu_compat_ops, dtype + ) + + output_dir.mkdir(parents=True, exist_ok=True) + pte = output_dir / "phonon.pte" + with pte.open("wb") as f: + program.write_to_file(f) + data_file = None + if backend == "cuda": + program.write_tensor_data_to_file(str(output_dir)) + data_file = "aoti_cuda_blob.ptd" + for name in ("tokenizer.json", "tokenizer.model", "default-voice.safetensors"): + if (checkpoint_dir / name).exists(): + shutil.copyfile(checkpoint_dir / name, output_dir / name) + + contract = { + "format_version": metadata["format_version"], + "backend": backend, + "dtype": dtype, + "quantize": quantize, + "qlinear": quantize_scheme(backend, quantize), + "webgpu_compat_ops": webgpu_compat_ops, + "files": { + "program": pte.name, + "data": data_file, + "tokenizer": "tokenizer.json", + "default_voice": "default-voice.safetensors", + }, + "constants": {k: v for k, v in metadata.items() if not torch.is_tensor(v)}, + "constant_tensors": {"bos_embeds": [1, 1, cfg.dim]}, + "dynamic": {"S": [1, cfg.max_prefill]}, + "methods": io_contract(cfg), + "mimi_state": [[n, list(s)] for n, s in mimi_state_spec(cfg)], + "delegation": delegation, + } + (output_dir / "phonon.json").write_text(json.dumps(contract, indent=2) + "\n") + print(json.dumps(delegation, indent=2)) + print(f"Saved {pte} ({pte.stat().st_size / 2**20:.1f} MiB)") + return pte + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=BACKENDS, default="xnnpack") + parser.add_argument("--output-dir", type=Path, default=Path("phonon_exports")) + parser.add_argument( + "--checkpoint-dir", + help="Local directory with model.safetensors, config.json, tokenizer.json " + "(default: download --hf-model at --revision)", + ) + parser.add_argument("--hf-model", default=HF_MODEL) + parser.add_argument("--revision", default=HF_REVISION) + parser.add_argument( + "--quantize", + choices=("4bit", "8bit"), + help="Quantize the flow LM's linear layers with the scheme of the backend " + "(see QUANTIZE_SCHEMES; default: none)", + ) + parser.add_argument("--quantize-group-size", type=int) + parser.add_argument( + "--dtype", + choices=tuple(DTYPES), + default="fp32", + help="Weight and matmul dtype; inputs, outputs and state stay float32", + ) + parser.add_argument( + "--webgpu-compat-ops", + action="store_true", + help="Replace layer_norm, elu, silu and conv_transpose1d with equivalent " + "ops for WebGPU runtimes built without those kernels", + ) + args = parser.parse_args() + # Before the download, so that an unsupported combination fails first. + try: + quantize_scheme(args.backend, args.quantize) + except ValueError as e: + parser.error(str(e)) + export_phonon( + resolve_checkpoint(args), + args.output_dir, + backend=args.backend, + quantize=args.quantize, + quantize_group_size=args.quantize_group_size, + webgpu_compat_ops=args.webgpu_compat_ops, + dtype=args.dtype, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/models/phonon/main.cpp b/examples/models/phonon/main.cpp new file mode 100644 index 00000000000..59a53af3026 --- /dev/null +++ b/examples/models/phonon/main.cpp @@ -0,0 +1,158 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Text to speech with an exported Phonon program: +// +// phonon_runner --model_path phonon_exports/xnnpack/phonon.pte \ +// --text "Hello from ExecuTorch." --output hello.wav + +#include "runner.h" + +#include +#include +#include +#include +#include +#include +#include + +#include + +DEFINE_string( + model_path, + "phonon_exports/xnnpack/phonon.pte", + "Exported program"); +DEFINE_string(data_path, "", "External CUDA data file (aoti_cuda_blob.ptd)"); +DEFINE_string( + tokenizer_path, + "", + "SentencePiece tokenizer.model (default: next to the program)"); +DEFINE_string( + voice_path, + "", + "Voice .safetensors (default: default-voice.safetensors next to the " + "program)"); +DEFINE_string(text, "", "Text to synthesize"); +DEFINE_string(output, "phonon.wav", "Output WAV file (16-bit PCM, mono)"); +DEFINE_double( + temperature, + -1, + "Flow noise variance; 0 removes the noise (default: the model's, 0.3)"); +DEFINE_uint64(seed, 0, "Noise seed"); + +namespace { + +std::string sibling(const std::string& path, const std::string& name) { + const auto slash = path.find_last_of("/\\"); + return slash == std::string::npos ? name : path.substr(0, slash + 1) + name; +} + +void write_u32(std::ofstream& out, uint32_t v) { + const unsigned char b[4] = { + static_cast(v), + static_cast(v >> 8), + static_cast(v >> 16), + static_cast(v >> 24)}; + out.write(reinterpret_cast(b), 4); +} + +void write_u16(std::ofstream& out, uint16_t v) { + const unsigned char b[2] = { + static_cast(v), static_cast(v >> 8)}; + out.write(reinterpret_cast(b), 2); +} + +// 16-bit mono PCM, rounded like run_phonon.py (round half to even). +void write_wav( + const std::string& path, + const std::vector& pcm, + int64_t sample_rate) { + std::ofstream out(path, std::ios::binary); + const auto data_bytes = static_cast(pcm.size() * 2); + out.write("RIFF", 4); + write_u32(out, 36 + data_bytes); + out.write("WAVEfmt ", 8); + write_u32(out, 16); + write_u16(out, 1); + write_u16(out, 1); + write_u32(out, static_cast(sample_rate)); + write_u32(out, static_cast(sample_rate * 2)); + write_u16(out, 2); + write_u16(out, 16); + out.write("data", 4); + write_u32(out, data_bytes); + for (float sample : pcm) { + double value = std::isfinite(sample) + ? std::nearbyint(static_cast(sample) * 32767.0) + : 0.0; + value = std::clamp(value, -32768.0, 32767.0); + write_u16(out, static_cast(static_cast(value))); + } + if (!out) { + throw std::runtime_error("Could not write " + path); + } +} + +double seconds_since(std::chrono::steady_clock::time_point start) { + return std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + try { + if (FLAGS_text.empty()) { + throw std::invalid_argument("Provide --text"); + } +#ifdef EXECUTORCH_BUILD_CUDA + if (FLAGS_data_path.empty()) { + throw std::invalid_argument( + "Provide --data_path pointing to aoti_cuda_blob.ptd for CUDA"); + } +#endif + auto start = std::chrono::steady_clock::now(); + phonon::Runner runner( + FLAGS_model_path, + FLAGS_tokenizer_path.empty() + ? sibling(FLAGS_model_path, "tokenizer.model") + : FLAGS_tokenizer_path, + FLAGS_data_path); + const auto voice = runner.load_voice( + FLAGS_voice_path.empty() + ? sibling(FLAGS_model_path, "default-voice.safetensors") + : FLAGS_voice_path); + const double load_s = seconds_since(start); + + start = std::chrono::steady_clock::now(); + const auto result = + runner.synthesize(FLAGS_text, voice, FLAGS_temperature, FLAGS_seed); + const double synth_s = seconds_since(start); + if (result.pcm.empty()) { + throw std::runtime_error("No audio generated"); + } + write_wav(FLAGS_output, result.pcm, runner.sample_rate()); + + const double audio_s = + static_cast(result.pcm.size()) / runner.sample_rate(); + std::printf( + "Wrote %s: %.2f s of audio, %lld frames, load %.2f s, synthesis " + "%.2f s (RTF %.2f)\n", + FLAGS_output.c_str(), + audio_s, + static_cast(result.frames), + load_s, + synth_s, + synth_s / audio_s); + } catch (const std::exception& error) { + std::cerr << "Phonon: " << error.what() << '\n'; + return 1; + } + return 0; +} diff --git a/examples/models/phonon/model.py b/examples/models/phonon/model.py new file mode 100644 index 00000000000..9701bf4d892 --- /dev/null +++ b/examples/models/phonon/model.py @@ -0,0 +1,1015 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""PyTorch definition of Gradium's Phonon text-to-speech model. + +Phonon (https://huggingface.co/gradium/phonon-31m-en) builds on Kyutai's +Pocket TTS and the Mimi codec. The modules below follow the official runtime +(xn-ptts, see NOTICE) and their parameter names match the keys of the official +`model.safetensors`, so the checkpoint loads with `strict=True`. + +Generation runs four programs that `export_phonon.py` exports as methods of one +`.pte` file (see README.md for the I/O contract): + + text_embed(tokens) text conditioner lookup, [1, S] -> [1, S, D] + voice_embed(voice) voice prompt projection, [1, T, D] -> [1, T, D] + flow_lm(embeds, input_pos, noise) + causal transformer with its KV cache inside the + program, then the EOS logit and the one-step + flow-matching head for the last position + mimi_decode(latent, *state) + one 80 ms frame of the streaming Mimi decoder; all + streaming state is explicit input and output +""" + +import json +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from executorch.examples.models.llama.rope import hf_apply_rotary_emb +from torch import nn + +# Checkpoint tensors the official runtime drops: a training-only loss weighting +# network. +TRAINING_ONLY_PREFIXES = ("flow_lm.flow.w_s_t.",) + +# Attention implementations for the flow LM KV cache (see FlowLMStep). +ATTENTION_IMPLS = ("custom_sdpa", "sdpa") + + +@dataclass(frozen=True) +class PhononConfig: + # Flow LM + dim: int = 512 + num_heads: int = 8 + num_layers: int = 6 + ffn_dim: int = 2048 + max_period: float = 10000.0 + text_bins: int = 4000 + latent_dim: int = 32 + flow_depth: int = 6 + eos_threshold: float = 4.0 + # Conditioners summed into every generated step + num_speakers_bins: int = 31 + num_speakers_dim: int = 16 + num_speakers_values: Tuple[str, ...] = ("0", "1", "2") + duration_dim: int = 128 + duration_scale: float = 1000.0 + duration_max_period: float = 10000.0 + # Mimi decoder + sample_rate: int = 24000 + frame_rate: float = 12.5 + mimi_dim: int = 512 + n_filters: int = 64 + ratios: Tuple[int, ...] = (6, 5, 4) + kernel_size: int = 7 + last_kernel_size: int = 3 + residual_kernel_size: int = 3 + compress: int = 2 + mimi_heads: int = 8 + mimi_layers: int = 2 + mimi_ffn_dim: int = 2048 + mimi_context: int = 250 + mimi_max_period: float = 10000.0 + # Export limits + max_context: int = 512 + max_prefill: int = 128 + # Generation defaults of the official runtime + temperature: float = 0.3 + + @property + def head_dim(self) -> int: + return self.dim // self.num_heads + + @property + def mimi_head_dim(self) -> int: + return self.mimi_dim // self.mimi_heads + + @property + def hop_length(self) -> int: + return math.prod(self.ratios) + + @property + def upsample_stride(self) -> int: + stride = self.sample_rate / self.hop_length / self.frame_rate + if stride != int(stride): + raise ValueError("sample_rate / prod(ratios) / frame_rate must be whole") + return int(stride) + + @property + def frame_samples(self) -> int: + return self.hop_length * self.upsample_stride + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "PhononConfig": + """Parse the official `config.json`, rejecting variants this example + does not implement.""" + lm, mimi = data["flow_lm"], data["mimi"] + checks = { + "lsd_decode_steps": (data.get("lsd_decode_steps"), 1), + "flow_lm.time_rms_norm": (lm.get("time_rms_norm"), "rms"), + "flow_lm.gelu": (lm.get("gelu"), "erf"), + "mimi.gelu": (mimi.get("gelu"), "erf"), + "mimi.channels": (mimi.get("channels"), 1), + "mimi.n_residual_layers": (mimi.get("n_residual_layers"), 1), + "flow_lm.lut_dim": (lm.get("lut_dim"), lm["d_model"]), + "flow_lm.flow_dim": (lm.get("flow_dim"), lm["d_model"]), + "fuser.sum": ( + data["fuser"]["sum"], + ["num_speakers", "duration_delta"], + ), + } + for name, (got, want) in checks.items(): + if got != want: + raise ValueError(f"Unsupported Phonon config: {name}={got!r}") + conds = {c["name"]: c for c in data["conditioners"]} + ns, dd = conds["num_speakers"]["lut"], conds["duration_delta"]["continuous"] + return cls( + dim=lm["d_model"], + num_heads=lm["num_heads"], + num_layers=lm["num_layers"], + ffn_dim=lm["dim_feedforward"], + max_period=lm["max_period"], + text_bins=lm["n_bins"], + latent_dim=lm["ldim"], + flow_depth=lm["flow_depth"], + eos_threshold=data["eos_threshold"], + num_speakers_bins=ns["n_bins"], + num_speakers_dim=ns["dim"], + num_speakers_values=tuple(ns["possible_values"]), + duration_dim=dd["dim"], + duration_scale=dd["scale_factor"], + duration_max_period=dd["max_period"], + sample_rate=mimi["sample_rate"], + frame_rate=mimi["frame_rate"], + mimi_dim=mimi["dimension"], + n_filters=mimi["n_filters"], + ratios=tuple(mimi["ratios"]), + kernel_size=mimi["kernel_size"], + last_kernel_size=mimi["last_kernel_size"], + residual_kernel_size=mimi["residual_kernel_size"], + compress=mimi["compress"], + mimi_heads=mimi["transformer_num_heads"], + mimi_layers=mimi["transformer_num_layers"], + mimi_ffn_dim=mimi["transformer_dim_feedforward"], + mimi_context=mimi["transformer_context"], + mimi_max_period=mimi["transformer_max_period"], + ) + + @classmethod + def from_json(cls, path: str | Path) -> "PhononConfig": + return cls.from_dict(json.loads(Path(path).read_text())) + + +# --------------------------------------------------------------------------- +# Building blocks (parameter names follow the checkpoint) +# --------------------------------------------------------------------------- + + +class Attention(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.in_proj = nn.Linear(dim, 3 * dim, bias=False) + self.out_proj = nn.Linear(dim, dim, bias=False) + + +class LayerScale(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.full((dim,), 0.01)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.scale * x + + +class TransformerLayer(nn.Module): + """Pre-norm transformer layer: attention (computed by the caller) and an + erf-GELU feed-forward block, with optional layer scale (Mimi).""" + + def __init__(self, dim: int, ffn_dim: int, layer_scale: bool): + super().__init__() + self.self_attn = Attention(dim) + self.norm1 = nn.LayerNorm(dim, eps=1e-5) + self.norm2 = nn.LayerNorm(dim, eps=1e-5) + self.linear1 = nn.Linear(dim, ffn_dim, bias=False) + self.linear2 = nn.Linear(ffn_dim, dim, bias=False) + self.gelu = nn.GELU() + if layer_scale: + self.layer_scale_1 = LayerScale(dim) + self.layer_scale_2 = LayerScale(dim) + else: + self.layer_scale_1 = nn.Identity() + self.layer_scale_2 = nn.Identity() + + def feed_forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.linear2(self.gelu(self.linear1(self.norm2(x)))) + return x + self.layer_scale_2(h) + + +class Transformer(nn.Module): + def __init__(self, n: int, dim: int, ffn_dim: int, layer_scale: bool): + super().__init__() + self.layers = nn.ModuleList( + TransformerLayer(dim, ffn_dim, layer_scale) for _ in range(n) + ) + + +def rope_inv_freq(head_dim: int, max_period: float) -> torch.Tensor: + half = head_dim // 2 + i = torch.arange(half, dtype=torch.float32) + return 1.0 / torch.pow(torch.tensor(max_period, dtype=torch.float32), i / half) + + +def rope_tables(n: int, head_dim: int, max_period: float): + """cos/sin [n, head_dim] for rotate-half RoPE at positions 0..n-1.""" + ang = torch.arange(n, dtype=torch.float32)[:, None] * rope_inv_freq( + head_dim, max_period + ) + cos, sin = torch.cos(ang), torch.sin(ang) + return torch.cat([cos, cos], -1), torch.cat([sin, sin], -1) + + +def interleaved_to_half_split(in_proj: torch.Tensor, num_heads: int) -> torch.Tensor: + """Reorder the q and k rows of a fused [3D, D] in_proj weight from + interleaved RoPE pairs (2i, 2i+1) to rotate-half pairs (i, i + head_dim/2). + + Phonon is trained with interleaved RoPE. Applying the same permutation to q + and k leaves every q.k product unchanged, and rotate-half RoPE on the new + layout equals interleaved RoPE on the old one, so the model's function is + unchanged while the graph uses the common rotate-half form that backends + fuse (for example Vulkan's apply_rotary_emb_hf). + """ + dim = in_proj.shape[1] + head_dim = dim // num_heads + per_head = torch.cat([torch.arange(0, head_dim, 2), torch.arange(1, head_dim, 2)]) + perm = torch.cat([h * head_dim + per_head for h in range(num_heads)]) + q, k, v = in_proj[:dim], in_proj[dim : 2 * dim], in_proj[2 * dim :] + return torch.cat([q[perm], k[perm], v], 0).contiguous() + + +# --------------------------------------------------------------------------- +# Flow LM +# --------------------------------------------------------------------------- + + +class TimeRMSNorm(nn.Module): + """Uncentred RMS norm with a learnt gain (config time_rms_norm: rms).""" + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.alpha = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt((x * x).mean(-1, keepdim=True) + self.eps) * self.alpha + + +class TimestepEmbedder(nn.Module): + def __init__(self, dim: int, freq_dim: int = 256): + super().__init__() + self.freqs = nn.Parameter(torch.zeros(freq_dim // 2)) + self.mlp = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim), + TimeRMSNorm(dim), + ) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + args = t * self.freqs + return self.mlp(torch.cat([torch.cos(args), torch.sin(args)], -1)) + + +class AdaLNResBlock(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + self.in_ln = nn.LayerNorm(dim, eps=1e-6) + self.mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 3 * dim)) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + d = self.dim + ada = self.adaLN_modulation(y) + shift, scale, gate = ada[..., :d], ada[..., d : 2 * d], ada[..., 2 * d :] + h = self.in_ln(x) * (scale + 1.0) + shift + return x + gate * self.mlp(h) + + +class FinalLayer(nn.Module): + def __init__(self, dim: int, out_dim: int): + super().__init__() + self.dim = dim + self.norm_final = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.linear = nn.Linear(dim, out_dim) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + ada = self.adaLN_modulation(y) + shift, scale = ada[..., : self.dim], ada[..., self.dim :] + return self.linear(self.norm_final(x) * (scale + 1.0) + shift) + + +class FlowNet(nn.Module): + """SimpleMLPAdaLN velocity network of the flow-matching head.""" + + def __init__(self, dim: int, latent_dim: int, depth: int): + super().__init__() + self.time_embed = nn.ModuleList([TimestepEmbedder(dim), TimestepEmbedder(dim)]) + self.cond_embed = nn.Linear(dim, dim) + self.input_proj = nn.Linear(latent_dim, dim) + self.res_blocks = nn.ModuleList(AdaLNResBlock(dim) for _ in range(depth)) + self.final_layer = FinalLayer(dim, latent_dim) + + def lsd_time_condition(self) -> torch.Tensor: + """Time condition of a single lagrangian self-distillation step from + s = 0 to t = 1: the mean of the two time embeddings, [1, dim].""" + s, t = torch.zeros(1, 1), torch.ones(1, 1) + return (self.time_embed[0](s) + self.time_embed[1](t)) / 2 + + def forward( + self, c: torch.Tensor, x: torch.Tensor, t_cond: torch.Tensor + ) -> torch.Tensor: + x = self.input_proj(x) + y = t_cond + self.cond_embed(c) + for block in self.res_blocks: + x = block(x, y) + return self.final_layer(x, y) + + +class Conditioner(nn.Module): + def __init__( + self, + dim: int, + embed: Optional[Tuple[int, int]] = None, + proj_in: Optional[int] = None, + ): + super().__init__() + # Unused at inference; kept so the checkpoint loads strictly. + self.learnt_padding = nn.Parameter(torch.zeros(1, 1, dim)) + if embed is not None: + self.embed = nn.Embedding(*embed) + if proj_in is not None: + self.output_proj = nn.Linear(proj_in, dim, bias=False) + + +class Conditioners(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + d = cfg.dim + self.duration_delta = Conditioner(d, proj_in=cfg.duration_dim) + self.num_speakers = Conditioner( + d, + embed=(cfg.num_speakers_bins + 1, cfg.num_speakers_dim), + proj_in=cfg.num_speakers_dim, + ) + self.speaker_wavs = Conditioner(d, proj_in=d) + self.transcript_in_segment = Conditioner(d, embed=(cfg.text_bins + 1, d)) + + +class ConditionProvider(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + self.conditioners = Conditioners(cfg) + + +def sinusoidal_embedding(position: float, dim: int, max_period: float) -> torch.Tensor: + half = dim // 2 + i = torch.arange(half, dtype=torch.float32) + phase = position / torch.pow( + torch.tensor(max_period, dtype=torch.float32), i / (half - 1) + ) + return torch.cat([torch.cos(phase), torch.sin(phase)]) + + +class FlowLM(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + self.cfg = cfg + d, ld = cfg.dim, cfg.latent_dim + self.bos_emb = nn.Parameter(torch.zeros(ld)) + self.emb_mean = nn.Parameter(torch.zeros(ld)) + self.emb_std = nn.Parameter(torch.ones(ld)) + self.condition_provider = ConditionProvider(cfg) + self.flow_net = FlowNet(d, ld, cfg.flow_depth) + self.input_linear = nn.Linear(ld, d, bias=False) + self.out_eos = nn.Linear(d, 1) + self.out_norm = nn.LayerNorm(d, eps=1e-5) + self.transformer = Transformer(cfg.num_layers, d, cfg.ffn_dim, False) + + @property + def conditioners(self) -> Conditioners: + return self.condition_provider.conditioners + + def summed_conditions( + self, num_speakers: str = "1", duration_delta: float = 0.0 + ) -> torch.Tensor: + """The conditions added to every generated step's input, [1, 1, D]. + Defaults are the official runtime's.""" + c, cfg = self.conditioners, self.cfg + idx = cfg.num_speakers_values.index(num_speakers) + ns = c.num_speakers.output_proj(c.num_speakers.embed.weight[idx : idx + 1]) + dd = c.duration_delta.output_proj( + sinusoidal_embedding( + cfg.duration_scale * duration_delta, + cfg.duration_dim, + cfg.duration_max_period, + )[None] + ) + return (ns + dd).view(1, 1, cfg.dim) + + def bos_embeds(self) -> torch.Tensor: + """Input embedding of the first generated step, [1, 1, D].""" + bos = self.input_linear(self.bos_emb.view(1, 1, -1)) + return bos + self.summed_conditions() + + +# --------------------------------------------------------------------------- +# Mimi decoder +# --------------------------------------------------------------------------- + + +class StreamingConv1d(nn.Module): + """Causal stride-1 conv; the last kernel_size - 1 input frames carry over.""" + + def __init__(self, cin: int, cout: int, kernel_size: int): + super().__init__() + self.conv = nn.Conv1d(cin, cout, kernel_size) + + @property + def state_len(self) -> int: + return self.conv.kernel_size[0] - 1 + + def forward( + self, x: torch.Tensor, prev: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if prev is None: + return self.conv(x), None + xx = torch.cat([prev, x], 2) + return self.conv(xx), xx[:, :, xx.shape[2] - prev.shape[2] :] + + +class StreamingConvTranspose1d(nn.Module): + """Transposed conv; the overlapping kernel_size - stride outputs are held + back (without bias) and added to the next frame's head.""" + + def __init__( + self, + cin: int, + cout: int, + kernel_size: int, + stride: int, + groups: int = 1, + bias: bool = True, + ): + super().__init__() + self.convtr = nn.ConvTranspose1d( + cin, cout, kernel_size, stride=stride, groups=groups, bias=bias + ) + + @property + def state_len(self) -> int: + return self.convtr.kernel_size[0] - self.convtr.stride[0] + + def forward( + self, x: torch.Tensor, partial: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + y = self.convtr(x) + pt, n = partial.shape[2], y.shape[2] + out = torch.cat([y[:, :, :pt] + partial, y[:, :, pt : n - pt]], 2) + tail = y[:, :, n - pt :] + bias = self.convtr.bias + if bias is not None: + tail = tail - bias.view(1, -1, 1) + return out, tail + + +class SEANetResBlock(nn.Module): + def __init__(self, dim: int, kernel_size: int, compress: int): + super().__init__() + hidden = dim // compress + self.block = nn.ModuleList( + [ + nn.ELU(), + StreamingConv1d(dim, hidden, kernel_size), + nn.ELU(), + StreamingConv1d(hidden, dim, 1), + ] + ) + + def forward(self, x: torch.Tensor, prev: torch.Tensor): + h, prev = self.block[1](self.block[0](x), prev) + h, _ = self.block[3](self.block[2](h), None) + return x + h, prev + + +class SEANetDecoder(nn.Module): + """Layout of `decoder.model`: conv, then (ELU, transposed conv, residual + block) per ratio, then ELU and the output conv.""" + + def __init__(self, cfg: PhononConfig): + super().__init__() + mult = 2 ** len(cfg.ratios) + layers: List[nn.Module] = [ + StreamingConv1d(cfg.mimi_dim, mult * cfg.n_filters, cfg.kernel_size) + ] + for ratio in cfg.ratios: + cin = mult * cfg.n_filters + layers += [ + nn.ELU(), + StreamingConvTranspose1d(cin, cin // 2, 2 * ratio, ratio), + SEANetResBlock(cin // 2, cfg.residual_kernel_size, cfg.compress), + ] + mult //= 2 + layers += [nn.ELU(), StreamingConv1d(cfg.n_filters, 1, cfg.last_kernel_size)] + self.model = nn.ModuleList(layers) + + +class MimiTransformer(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + self.transformer = Transformer( + cfg.mimi_layers, cfg.mimi_dim, cfg.mimi_ffn_dim, True + ) + + +class Upsample(nn.Module): + """Depthwise transposed conv from the codec frame rate to the transformer + rate.""" + + def __init__(self, dim: int, stride: int): + super().__init__() + self.convtr = StreamingConvTranspose1d( + dim, dim, 2 * stride, stride, groups=dim, bias=False + ) + + def forward(self, x: torch.Tensor, partial: torch.Tensor): + return self.convtr(x, partial) + + +class Quantizer(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + self.output_proj = nn.Conv1d(cfg.latent_dim, cfg.mimi_dim, 1, bias=False) + + +class MimiDecoder(nn.Module): + def __init__(self, cfg: PhononConfig): + super().__init__() + s = cfg.upsample_stride + self.quantizer = Quantizer(cfg) + self.upsample = Upsample(cfg.mimi_dim, s) + self.decoder_transformer = MimiTransformer(cfg) + self.decoder = SEANetDecoder(cfg) + + +class Phonon(nn.Module): + def __init__(self, cfg: Optional[PhononConfig] = None): + super().__init__() + cfg = cfg if cfg is not None else PhononConfig() + self.cfg = cfg + self.flow_lm = FlowLM(cfg) + self.mimi = MimiDecoder(cfg) + # Layout of the q/k rows of every in_proj (see interleaved_to_half_split). + self.rope_layout = "interleaved" + + @classmethod + def from_checkpoint( + cls, checkpoint: str | Path, cfg: Optional[PhononConfig] = None + ) -> "Phonon": + """Load the official `model.safetensors` (bf16 flow LM, fp32 Mimi) in + fp32. Every tensor is loaded with strict=True except the training-only + tensors in TRAINING_ONLY_PREFIXES, which the official runtime drops too. + """ + from safetensors.torch import load_file + + checkpoint = Path(checkpoint) + if cfg is None: + cfg = PhononConfig.from_json(checkpoint.parent / "config.json") + state = { + k: v.float() + for k, v in load_file(str(checkpoint)).items() + if not k.startswith(TRAINING_ONLY_PREFIXES) + } + model = cls(cfg) + model.load_state_dict(state, strict=True) + return model.eval().requires_grad_(False).to_half_split_rope() + + def to_half_split_rope(self) -> "Phonon": + """Permute every in_proj once for rotate-half RoPE (exact, in place).""" + if self.rope_layout == "half_split": + return self + blocks = [ + (self.flow_lm.transformer, self.cfg.num_heads), + (self.mimi.decoder_transformer.transformer, self.cfg.mimi_heads), + ] + with torch.no_grad(): + for transformer, heads in blocks: + for layer in transformer.layers: + w = layer.self_attn.in_proj.weight + w.copy_(interleaved_to_half_split(w, heads)) + self.rope_layout = "half_split" + return self + + +def load_voice(path: str | Path) -> torch.Tensor: + """Speaker latents [1, T, D] from a Phonon voice `.safetensors` file + (stored as speaker_wavs [1, D, T]).""" + from safetensors.torch import load_file + + tensors = load_file(str(path)) + if list(tensors) != ["speaker_wavs"]: + raise ValueError(f"{path}: expected one tensor named speaker_wavs") + return tensors["speaker_wavs"].float().transpose(1, 2).contiguous() + + +# --------------------------------------------------------------------------- +# Export programs +# --------------------------------------------------------------------------- + + +class TextEmbedding(nn.Module): + """tokens i64[1, S] -> embeds f32[1, S, D] (text conditioner lookup).""" + + def __init__(self, model: Phonon): + super().__init__() + self.embed = model.flow_lm.conditioners.transcript_in_segment.embed + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + return self.embed(tokens) + + +class VoiceEmbedding(nn.Module): + """voice f32[1, T, D] -> embeds f32[1, T, D] (voice prompt projection). + + Taking the speaker latents time-major keeps the dynamic length out of any + transpose, which not every delegate resizes.""" + + def __init__(self, model: Phonon): + super().__init__() + self.proj = model.flow_lm.conditioners.speaker_wavs.output_proj + + def forward(self, voice: torch.Tensor) -> torch.Tensor: + d = voice.shape[2] + return self.proj(voice.reshape(-1, d)).reshape(1, -1, d) + + +class FlowLMStep(nn.Module): + """forward(embeds f32[1, S, D], input_pos i64[1], noise f32[1, latent_dim]) + -> latent f32[1, latent_dim], eos_logit f32[1, 1], next_embeds f32[1, 1, D] + + Writes KV cache positions [input_pos, input_pos + S) and attends causally + over [0, input_pos + S). The head runs on the last position: + eos_logit = out_eos(h), latent = noise + flow_net(h, noise) (one flow step) + and next_embeds = input_linear(latent) + conditions, the next step's input. + Prompt calls (voice, then text) ignore the outputs. + + attention selects how the KV cache is stored and attended: + custom_sdpa torch.ops.llama.sdpa_with_kv_cache (ExecuTorch CPU custom op; + Vulkan and WebGPU lower it to a fused kernel) + sdpa index_copy_ into the cache + F.scaled_dot_product_attention + with an explicit causal mask (AOTInductor, MLX) + """ + + def __init__(self, model: Phonon, attention: str = "custom_sdpa"): + super().__init__() + if attention not in ATTENTION_IMPLS: + raise ValueError(f"attention must be one of {ATTENTION_IMPLS}") + if model.rope_layout != "half_split": + raise ValueError("Call Phonon.to_half_split_rope() first") + cfg, lm = model.cfg, model.flow_lm + self.cfg, self.attention = cfg, attention + self.layers = lm.transformer.layers + self.out_norm, self.out_eos = lm.out_norm, lm.out_eos + self.input_linear, self.flow_net = lm.input_linear, lm.flow_net + with torch.no_grad(): + self.register_buffer("cond", lm.summed_conditions().clone()) + self.register_buffer("t_cond", lm.flow_net.lsd_time_condition().clone()) + cos, sin = rope_tables(cfg.max_context, cfg.head_dim, cfg.max_period) + self.register_buffer("rope_cos", cos) + self.register_buffer("rope_sin", sin) + if attention == "custom_sdpa": + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + shape = (1, cfg.max_context, cfg.num_heads, cfg.head_dim) + else: + shape = (1, cfg.num_heads, cfg.max_context, cfg.head_dim) + for i in range(cfg.num_layers): + self.register_buffer(f"k_cache_{i}", torch.zeros(shape), persistent=False) + self.register_buffer(f"v_cache_{i}", torch.zeros(shape), persistent=False) + + def reset(self) -> None: + for name, buf in self.named_buffers(): + if "_cache_" in name: + buf.zero_() + + def _attend_custom(self, i, q, k, v, input_pos): + cfg = self.cfg + seq = q.shape[1] + start = input_pos[0].item() + torch._check(start >= 0) + torch._check(start + seq <= cfg.max_context) + cos = self.rope_cos.narrow(0, start, seq) + sin = self.rope_sin.narrow(0, start, seq) + q, k = hf_apply_rotary_emb(q, k, cos, sin) + return torch.ops.llama.sdpa_with_kv_cache( + q, + k, + v, + getattr(self, f"k_cache_{i}"), + getattr(self, f"v_cache_{i}"), + start, + seq, + None, + 0.0, + True, + None, + ) + + def _attend_sdpa(self, i, q, k, v, input_pos): + cfg = self.cfg + seq = q.shape[1] + pos = input_pos + torch.arange(seq, dtype=torch.long) + q, k = hf_apply_rotary_emb(q, k, self.rope_cos[pos], self.rope_sin[pos]) + k_cache, v_cache = getattr(self, f"k_cache_{i}"), getattr(self, f"v_cache_{i}") + k_cache.index_copy_(2, pos, k.transpose(1, 2)) + v_cache.index_copy_(2, pos, v.transpose(1, 2)) + keys = torch.arange(cfg.max_context, dtype=torch.long) + mask = torch.where(keys[None, :] <= pos[:, None], 0.0, float("-inf")) + o = F.scaled_dot_product_attention( + q.transpose(1, 2), k_cache, v_cache, attn_mask=mask + ) + return o.transpose(1, 2) + + def forward( + self, embeds: torch.Tensor, input_pos: torch.Tensor, noise: torch.Tensor + ): + cfg = self.cfg + seq = embeds.shape[1] + heads, hd, d = cfg.num_heads, cfg.head_dim, cfg.dim + attend = ( + self._attend_custom + if self.attention == "custom_sdpa" + else self._attend_sdpa + ) + x = embeds + for i, layer in enumerate(self.layers): + qkv = layer.self_attn.in_proj(layer.norm1(x)) + q = qkv[..., :d].reshape(1, seq, heads, hd) + k = qkv[..., d : 2 * d].reshape(1, seq, heads, hd) + v = qkv[..., 2 * d :].reshape(1, seq, heads, hd) + o = attend(i, q, k, v, input_pos) + x = x + layer.self_attn.out_proj(o.reshape(1, seq, d)) + x = layer.feed_forward(x) + h = self.out_norm(x[:, -1, :]) + eos_logit = self.out_eos(h) + latent = noise + self.flow_net(h, noise, self.t_cond) + next_embeds = self.input_linear(latent).reshape(1, 1, d) + self.cond + return latent, eos_logit, next_embeds + + +def mimi_state_spec(cfg: PhononConfig) -> List[Tuple[str, Tuple[int, ...]]]: + """Names and shapes of the streaming Mimi decoder state, in call order.""" + h, w, hd = cfg.mimi_heads, cfg.mimi_context, cfg.mimi_head_dim + spec: List[Tuple[str, Tuple[int, ...]]] = [ + ("upsample_partial", (1, cfg.mimi_dim, cfg.upsample_stride)) + ] + for i in range(cfg.mimi_layers): + spec += [(f"k{i}", (1, h, w, hd)), (f"v{i}", (1, h, w, hd))] + spec += [("kv_seen", (1, 1, 1, w)), ("rope_offset", (1,))] + mult = 2 ** len(cfg.ratios) + spec.append(("init_conv_prev", (1, cfg.mimi_dim, cfg.kernel_size - 1))) + for j, ratio in enumerate(cfg.ratios, 1): + c = mult * cfg.n_filters // 2 + spec += [ + (f"up{j}_partial", (1, c, ratio)), + (f"res{j}_prev", (1, c, cfg.residual_kernel_size - 1)), + ] + mult //= 2 + spec.append(("final_conv_prev", (1, cfg.n_filters, cfg.last_kernel_size - 1))) + return spec + + +def mimi_zero_state(cfg: PhononConfig) -> List[torch.Tensor]: + return [torch.zeros(shape) for _, shape in mimi_state_spec(cfg)] + + +class MimiDecodeStep(nn.Module): + """forward(latent f32[1, 1, latent_dim], *state) -> pcm f32[1, 1, frame], + *new_state, with state as listed by mimi_state_spec (zeros to start). + + One latent frame per call: denormalize, project, upsample x16, a 2-layer + transformer over a sliding window of the last mimi_context positions, and + the SEANet decoder. Every piece of streaming state (conv tails, transposed + conv partial sums, the window's keys and values, which window slots are + filled, and the RoPE offset) is an explicit input and output, so no backend + has to keep mutable state alive between calls. + """ + + def __init__(self, model: Phonon): + super().__init__() + if model.rope_layout != "half_split": + raise ValueError("Call Phonon.to_half_split_rope() first") + cfg, mimi = model.cfg, model.mimi + self.cfg = cfg + self.emb_std = model.flow_lm.emb_std + self.emb_mean = model.flow_lm.emb_mean + self.quantizer = mimi.quantizer + self.upsample = mimi.upsample + # Not `layers`: the CUDA backend shares weights between the methods of a + # program by name, and FlowLMStep.layers holds other weights. + self.transformer_layers = mimi.decoder_transformer.transformer.layers + self.decoder = mimi.decoder.model + t, w = cfg.upsample_stride, cfg.mimi_context + self.register_buffer( + "inv_freq", rope_inv_freq(cfg.mimi_head_dim, cfg.mimi_max_period)[None] + ) + self.register_buffer("steps", torch.arange(t, dtype=torch.float32)[:, None]) + i, j = torch.arange(t)[:, None], torch.arange(w + t)[None, :] + # Query i sits at window index w + i and sees keys i..w + i. + band = torch.where((j >= i) & (j <= i + w), 0.0, -1.0e30) + self.register_buffer("band", band) + self.register_buffer("ones_seen", torch.ones(1, 1, 1, t)) + + def _attention(self, layer, x, k_state, v_state, seen_all, cos, sin): + cfg = self.cfg + t, w, h, hd, d = ( + cfg.upsample_stride, + cfg.mimi_context, + cfg.mimi_heads, + cfg.mimi_head_dim, + cfg.mimi_dim, + ) + qkv = layer.self_attn.in_proj(layer.norm1(x)) # [t, 3d] + # [1, t, h, hd] is the layout of the rotate-half RoPE pattern. + q = qkv[:, :d].reshape(1, t, h, hd) + k = qkv[:, d : 2 * d].reshape(1, t, h, hd) + v = qkv[:, 2 * d :].reshape(t, h, hd) + q, k = hf_apply_rotary_emb(q, k, cos, sin) + q, k = q.reshape(t, h, hd), k.reshape(t, h, hd) + keys = torch.cat([k_state.reshape(h, w, hd), k.permute(1, 0, 2)], 1) + values = torch.cat([v_state.reshape(h, w, hd), v.permute(1, 0, 2)], 1) + mask = self.band + (seen_all.reshape(1, w + t) - 1.0) * 1.0e30 + scores = torch.bmm(q.permute(1, 0, 2), keys.transpose(1, 2)) * (hd**-0.5) + probs = torch.softmax(scores + mask, -1) + o = torch.bmm(probs, values).permute(1, 0, 2).reshape(t, d) + x = x + layer.layer_scale_1(layer.self_attn.out_proj(o)) + x = layer.feed_forward(x) + return ( + x, + keys[:, t:].reshape(1, h, w, hd), + values[:, t:].reshape(1, h, w, hd), + ) + + def forward(self, latent: torch.Tensor, *state: torch.Tensor): + cfg = self.cfg + n_layers, t = cfg.mimi_layers, cfg.upsample_stride + state = list(state) + up_partial = state[0] + kv = state[1 : 1 + 2 * n_layers] + seen, offset = state[1 + 2 * n_layers], state[2 + 2 * n_layers] + convs = state[3 + 2 * n_layers :] + + x = latent * self.emb_std + self.emb_mean # [1, 1, latent] + x = self.quantizer.output_proj(x.transpose(1, 2)) # [1, d, 1] + x, up_partial = self.upsample(x, up_partial) # [1, d, t] + + ang = (self.steps + offset) * self.inv_freq # [t, hd / 2] + cos = torch.cat([torch.cos(ang), torch.cos(ang)], -1) + sin = torch.cat([torch.sin(ang), torch.sin(ang)], -1) + seen_all = torch.cat([seen, self.ones_seen], 3) + x = x.reshape(cfg.mimi_dim, t).transpose(0, 1) # [t, d] + new_kv = [] + for i, layer in enumerate(self.transformer_layers): + x, k, v = self._attention( + layer, x, kv[2 * i], kv[2 * i + 1], seen_all, cos, sin + ) + new_kv += [k, v] + x = x.transpose(0, 1).reshape(1, cfg.mimi_dim, t) + + m = self.decoder + new_convs = [] + x, c = m[0](x, convs[0]) + new_convs.append(c) + for j in range(len(cfg.ratios)): + elu, up, res = m[1 + 3 * j], m[2 + 3 * j], m[3 + 3 * j] + x, p = up(elu(x), convs[1 + 2 * j]) + x, r = res(x, convs[2 + 2 * j]) + new_convs += [p, r] + pcm, c = m[-1](m[-2](x), convs[-1]) + new_convs.append(c) + return ( + pcm, + up_partial, + *new_kv, + seen_all[..., t:], + offset + float(t), + *new_convs, + ) + + +def example_inputs(cfg: PhononConfig) -> dict[str, Tuple[torch.Tensor, ...]]: + g = torch.Generator().manual_seed(0) + seq = 5 + return { + "text_embed": (torch.zeros(1, seq, dtype=torch.long),), + "voice_embed": (torch.randn(1, seq, cfg.dim, generator=g),), + "flow_lm": ( + torch.randn(1, seq, cfg.dim, generator=g), + torch.tensor([0], dtype=torch.long), + torch.zeros(1, cfg.latent_dim), + ), + "mimi_decode": ( + torch.randn(1, 1, cfg.latent_dim, generator=g), + *mimi_zero_state(cfg), + ), + } + + +def dynamic_shapes(cfg: PhononConfig) -> dict[str, Any]: + seq = torch.export.Dim("seq", min=1, max=cfg.max_prefill) + return { + "text_embed": ({1: seq},), + "voice_embed": ({1: seq},), + "flow_lm": ({1: seq}, None, None), + "mimi_decode": None, + } + + +def build_programs( + model: Phonon, attention: str = "custom_sdpa" +) -> dict[str, nn.Module]: + return { + "text_embed": TextEmbedding(model).eval(), + "voice_embed": VoiceEmbedding(model).eval(), + "flow_lm": FlowLMStep(model, attention).eval(), + "mimi_decode": MimiDecodeStep(model).eval(), + } + + +def program_metadata(model: Phonon) -> dict[str, Any]: + """Constant methods stored in the .pte next to the four programs.""" + cfg = model.cfg + with torch.no_grad(): + bos = model.flow_lm.bos_embeds().contiguous() + return { + "format_version": 1, + "sample_rate": cfg.sample_rate, + "frame_samples": cfg.frame_samples, + "frame_rate": cfg.frame_rate, + "dim": cfg.dim, + "latent_dim": cfg.latent_dim, + "max_context": cfg.max_context, + "max_prefill": cfg.max_prefill, + "eos_threshold": cfg.eos_threshold, + "temperature": cfg.temperature, + "num_mimi_state": len(mimi_state_spec(cfg)), + "bos_embeds": bos, + } + + +def io_contract(cfg: PhononConfig) -> dict[str, Any]: + """Inputs and outputs of every program; "S" is the dynamic sequence length + (1 <= S <= max_prefill). Identical for every backend.""" + d, ld = cfg.dim, cfg.latent_dim + state = [ + {"name": name, "shape": list(shape), "dtype": "float32"} + for name, shape in mimi_state_spec(cfg) + ] + + def t(name, shape, dtype="float32"): + return {"name": name, "shape": shape, "dtype": dtype} + + return { + "text_embed": { + "inputs": [t("tokens", [1, "S"], "int64")], + "outputs": [t("embeds", [1, "S", d])], + }, + "voice_embed": { + "inputs": [t("voice", [1, "S", d])], + "outputs": [t("embeds", [1, "S", d])], + }, + "flow_lm": { + "inputs": [ + t("embeds", [1, "S", d]), + t("input_pos", [1], "int64"), + t("noise", [1, ld]), + ], + "outputs": [ + t("latent", [1, ld]), + t("eos_logit", [1, 1]), + t("next_embeds", [1, 1, d]), + ], + }, + "mimi_decode": { + "inputs": [t("latent", [1, 1, ld])] + state, + "outputs": [t("pcm", [1, 1, cfg.frame_samples])] + state, + }, + } diff --git a/examples/models/phonon/run_phonon.py b/examples/models/phonon/run_phonon.py new file mode 100644 index 00000000000..b35e00cbcd2 --- /dev/null +++ b/examples/models/phonon/run_phonon.py @@ -0,0 +1,307 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Synthesize speech with an exported Phonon program (ExecuTorch pybindings). + + python -m executorch.examples.models.phonon.run_phonon \\ + --model phonon_exports/xnnpack/phonon.pte \\ + --text "Hello from ExecuTorch." --output hello.wav + +The host side is small and backend independent: basic text preparation and +sentence chunking, tokenization with tokenizer.json, the noise of the flow +head, the end-of-speech rule, and feeding next_embeds and the Mimi state back. +""" + +import argparse +import math +import re +import time +import wave +from pathlib import Path +from typing import Callable, List, Optional, Sequence, Tuple + +import numpy as np +import torch + +MAX_CHUNK_TOKENS = 50 # the official runtime's sentence chunk size + + +def prepare_text(text: str) -> Tuple[str, int]: + """Basic text preparation of the official runtime: collapse whitespace, + capitalize the first letter, end with a period if the text ends with a + letter or digit. Returns the text and how many frames to keep generating + after the end-of-speech logit fires (3 for texts of up to 4 words, else 1). + + Letters and digits are ASCII, as in the C++ runner. Numbers, symbols, + abbreviations, URLs and other non-word input are not normalized (the + official runtime has English rules for them).""" + words = text.split() + text = " ".join(words) + if not text: + return text, 3 + frames_after_eos = 3 if len(words) <= 4 else 1 + if "a" <= text[0] <= "z": + text = text[0].upper() + text[1:] + if text[-1].isascii() and text[-1].isalnum(): + text += "." + return text, frames_after_eos + + +def split_sentences(text: str, count_tokens: Callable[[str], int], max_tokens: int): + """Greedily pack whole sentences into chunks of at most max_tokens tokens; + a longer sentence is split at word boundaries.""" + chunks: List[str] = [] + current = "" + for sentence in re.split(r"(?<=[.!?])\s+", " ".join(text.split())): + pieces = [sentence] + if count_tokens(sentence) > max_tokens: + pieces = [] + piece = "" + for word in sentence.split(): + candidate = f"{piece} {word}".strip() + if piece and count_tokens(candidate) > max_tokens: + pieces.append(piece) + candidate = word + piece = candidate + pieces.append(piece) + for piece in pieces: + candidate = f"{current} {piece}".strip() + if current and count_tokens(candidate) > max_tokens: + chunks.append(current) + candidate = piece + current = candidate + if current: + chunks.append(current) + return chunks + + +def frame_budget(num_tokens: int, frame_rate: float) -> int: + return math.ceil((num_tokens / 3.0 + 2.0) * frame_rate) + + +class NoiseGenerator: + """Gaussian noise of the flow head, identical to phonon::NoiseGenerator in + the C++ runner: splitmix64 uniforms and the Box-Muller transform.""" + + MASK = (1 << 64) - 1 + + def __init__(self, seed: int): + self.state = seed & self.MASK + + def _uniform(self) -> float: + self.state = (self.state + 0x9E3779B97F4A7C15) & self.MASK + z = self.state + z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & self.MASK + z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & self.MASK + z ^= z >> 31 + return (z >> 11) * 2.0**-53 + + def normal(self, count: int, stddev: float) -> torch.Tensor: + out = [] + while len(out) < count: + radius = math.sqrt(-2.0 * math.log(1.0 - self._uniform())) + angle = 2.0 * math.pi * self._uniform() + out += [ + radius * math.cos(angle) * stddev, + radius * math.sin(angle) * stddev, + ] + return torch.tensor(out[:count], dtype=torch.float32) + + +class ExecuTorchPrograms: + """Methods of an exported phonon.pte, run with the ExecuTorch pybindings.""" + + def __init__(self, model_path: str, data_path: Optional[str] = None): + # Kernels for llama::sdpa_with_kv_cache (CPU backends) and quantized ops. + import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + from executorch.extension.pybindings.portable_lib import _load_for_executorch + + try: + import executorch.kernels.quantized # noqa: F401 + except ImportError: + pass + self.module = _load_for_executorch(model_path, data_path=data_path) + + def __call__(self, method: str, inputs: Sequence[torch.Tensor]) -> List: + return list(self.module.run_method(method, [x.contiguous() for x in inputs])) + + def constant(self, name: str): + return self.module.run_method(name, [])[0] + + def mimi_zero_state(self) -> List[torch.Tensor]: + meta = self.module.method_meta("mimi_decode") + return [ + torch.zeros(tuple(meta.input_tensor_meta(i).sizes())) + for i in range(1, meta.num_inputs()) + ] + + +class EagerPrograms: + """The same methods as eager PyTorch modules (for parity checks).""" + + def __init__(self, programs: dict, metadata: dict): + self.programs, self.metadata = programs, metadata + + def __call__(self, method: str, inputs: Sequence[torch.Tensor]) -> List: + with torch.no_grad(): + out = self.programs[method](*inputs) + return [t.clone() for t in (out if isinstance(out, tuple) else (out,))] + + def constant(self, name: str): + return self.metadata[name] + + def mimi_zero_state(self) -> List[torch.Tensor]: + from executorch.examples.models.phonon.model import mimi_zero_state + + return mimi_zero_state(self.programs["mimi_decode"].cfg) + + +def load_tokenizer(path: str | Path) -> Callable[[str], List[int]]: + """Encoder for the upstream tokenizer.json (no special tokens).""" + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(str(path)) + return lambda text: tokenizer.encode(text, add_special_tokens=False).ids + + +class PhononRunner: + def __init__(self, programs, tokenize: Callable[[str], List[int]]): + self.run = programs + self.tokenize = tokenize + c = programs.constant + if c("format_version") != 1: + raise ValueError("Unsupported Phonon program; re-export it") + self.sample_rate = c("sample_rate") + self.frame_samples = c("frame_samples") + self.frame_rate = c("frame_rate") + self.latent_dim = c("latent_dim") + self.max_context = c("max_context") + self.max_prefill = c("max_prefill") + self.eos_threshold = c("eos_threshold") + self.default_temperature = c("temperature") + self.bos_embeds = c("bos_embeds") + + def _flow(self, embeds, pos, noise): + latent, eos, next_embeds = self.run( + "flow_lm", [embeds, torch.tensor([pos], dtype=torch.long), noise] + ) + return latent, float(eos.reshape(-1)[0]), next_embeds + + def synthesize( + self, + text: str, + voice: torch.Tensor, + temperature: Optional[float] = None, + seed: int = 0, + noise_fn: Optional[Callable[[int], torch.Tensor]] = None, + ) -> Tuple[np.ndarray, dict]: + """Returns float PCM at sample_rate and generation info. noise_fn(k) + overrides the noise of step k (default: N(0, temperature) from + NoiseGenerator(seed)).""" + temperature = self.default_temperature if temperature is None else temperature + if noise_fn is None: + gen, scale = NoiseGenerator(seed), math.sqrt(temperature) + + def noise_fn(k): + return gen.normal(self.latent_dim, scale).reshape(1, -1) + + voice_embeds = self.run("voice_embed", [voice])[0] + voice_len = voice_embeds.shape[1] + if voice_len > self.max_prefill: + raise ValueError("Voice prompt is longer than max_prefill") + self._flow(voice_embeds, 0, torch.zeros(1, self.latent_dim)) + mimi_state = self.run.mimi_zero_state() + prepared, _ = prepare_text(text) + chunks = split_sentences( + prepared, lambda s: len(self.tokenize(s)), MAX_CHUNK_TOKENS + ) + pcm, info, step = [], {"chunks": [], "latents": [], "eos_logits": []}, 0 + for chunk in chunks: + chunk, frames_after_eos = prepare_text(chunk) + tokens = self.tokenize(chunk) + if not 0 < len(tokens) <= self.max_prefill: + raise ValueError(f"Chunk has an unsupported token count: {chunk}") + pos = voice_len + text_embeds = self.run("text_embed", [torch.tensor([tokens])])[0] + self._flow(text_embeds, pos, torch.zeros(1, self.latent_dim)) + pos += len(tokens) + budget = min( + frame_budget(len(tokens), self.frame_rate), self.max_context - pos + ) + embeds, countdown, frames = self.bos_embeds, None, 0 + for _ in range(budget): + latent, eos_logit, embeds = self._flow(embeds, pos, noise_fn(step)) + pos, step, frames = pos + 1, step + 1, frames + 1 + out = self.run("mimi_decode", [latent.reshape(1, 1, -1), *mimi_state]) + pcm.append(out[0].reshape(-1).numpy().copy()) + mimi_state = out[1:] + info["latents"].append(latent.reshape(-1).numpy().copy()) + info["eos_logits"].append(eos_logit) + if eos_logit > self.eos_threshold and countdown is None: + countdown = frames_after_eos + if countdown is not None: + if countdown == 0: + break + countdown -= 1 + info["chunks"].append({"text": chunk, "tokens": tokens, "frames": frames}) + info["frames"] = step + return np.concatenate(pcm) if pcm else np.zeros(0, np.float32), info + + +def write_wav(path: str | Path, pcm: np.ndarray, sample_rate: int) -> None: + pcm16 = np.clip(np.round(pcm.astype(np.float64) * 32767), -32768, 32767) + with wave.open(str(path), "wb") as f: + f.setnchannels(1) + f.setsampwidth(2) + f.setframerate(sample_rate) + f.writeframes(pcm16.astype(" +#include +#include +#include +#include + +#include +#include + +namespace phonon { +namespace { +using executorch::aten::ScalarType; +using executorch::aten::SizesType; +using executorch::extension::from_blob; +using executorch::extension::Module; +using executorch::runtime::Error; +using executorch::runtime::EValue; + +// The official runtime's sentence chunk size. +constexpr size_t kMaxChunkTokens = 50; +constexpr int kFormatVersion = 1; + +std::vector execute( + Module& model, + const std::string& method, + const std::vector& inputs = {}) { + auto result = model.execute(method, inputs); + if (!result.ok()) { + throw std::runtime_error( + method + " failed: " + executorch::runtime::to_string(result.error())); + } + return std::move(result.get()); +} + +const float* float_data(const EValue& value) { + const auto& tensor = value.toTensor(); + if (tensor.scalar_type() != ScalarType::Float) { + throw std::runtime_error("Expected a float32 model output"); + } + return tensor.const_data_ptr(); +} + +std::vector copy_floats(const EValue& value) { + const float* data = float_data(value); + return {data, data + value.toTensor().numel()}; +} + +// Code points of Python's str.isspace(), so that whitespace handling matches +// str.split() in run_phonon.py. +bool is_space(uint32_t c) { + return (c >= 0x09 && c <= 0x0d) || (c >= 0x1c && c <= 0x20) || c == 0x85 || + c == 0xa0 || c == 0x1680 || (c >= 0x2000 && c <= 0x200a) || c == 0x2028 || + c == 0x2029 || c == 0x202f || c == 0x205f || c == 0x3000; +} + +// Decodes the UTF-8 code point at text[i] and advances i; invalid bytes are +// returned as single code units. +uint32_t next_code_point(const std::string& text, size_t& i) { + const auto byte = [&](size_t k) { + return static_cast(text[k]); + }; + const unsigned char lead = byte(i); + const int length = lead < 0x80 ? 1 + : (lead >> 5) == 0x6 ? 2 + : (lead >> 4) == 0xe ? 3 + : (lead >> 3) == 0x1e ? 4 + : 1; + if (length == 1 || i + length > text.size()) { + ++i; + return lead; + } + uint32_t c = lead & (0x7f >> length); + for (int k = 1; k < length; ++k) { + if ((byte(i + k) & 0xc0) != 0x80) { + ++i; + return lead; + } + c = (c << 6) | (byte(i + k) & 0x3f); + } + i += length; + return c; +} + +std::vector split_words(const std::string& text) { + std::vector words; + std::string word; + for (size_t i = 0; i < text.size();) { + const size_t start = i; + if (is_space(next_code_point(text, i))) { + if (!word.empty()) { + words.push_back(std::move(word)); + word.clear(); + } + } else { + word.append(text, start, i - start); + } + } + if (!word.empty()) { + words.push_back(std::move(word)); + } + return words; +} + +std::string join(const std::vector& words) { + std::string out; + for (const auto& word : words) { + if (!out.empty()) { + out += ' '; + } + out += word; + } + return out; +} + +std::string join(const std::string& a, const std::string& b) { + return a.empty() ? b : b.empty() ? a : a + " " + b; +} + +bool is_ascii_alnum(char c) { + return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'z') || + (c >= 'A' && c <= 'Z'); +} + +std::vector split_at_sentence_ends(const std::string& text) { + std::vector sentences; + size_t start = 0; + for (size_t i = 1; i < text.size(); ++i) { + if (text[i] == ' ' && + (text[i - 1] == '.' || text[i - 1] == '!' || text[i - 1] == '?')) { + sentences.push_back(text.substr(start, i - start)); + start = i + 1; + } + } + sentences.push_back(text.substr(start)); + return sentences; +} + +std::vector split_long_sentence( + const std::string& sentence, + const std::function& count_tokens, + size_t max_tokens) { + std::vector pieces; + std::string piece; + for (const auto& word : split_words(sentence)) { + std::string candidate = join(piece, word); + if (!piece.empty() && count_tokens(candidate) > max_tokens) { + pieces.push_back(piece); + candidate = word; + } + piece = std::move(candidate); + } + pieces.push_back(piece); + return pieces; +} + +int64_t frame_budget(size_t num_tokens, double frame_rate) { + return static_cast( + std::ceil((static_cast(num_tokens) / 3.0 + 2.0) * frame_rate)); +} + +uint16_t read_u16_le(const unsigned char* p) { + return static_cast(p[0] | (p[1] << 8)); +} + +float half_to_float(uint16_t h) { + const uint32_t sign = (h & 0x8000u) << 16; + const int exponent = (h >> 10) & 0x1f; + const uint32_t mantissa = h & 0x3ffu; + float magnitude; + if (exponent == 0) { + magnitude = std::ldexp(static_cast(mantissa), -24); + } else if (exponent == 31) { + magnitude = mantissa ? NAN : INFINITY; + } else { + magnitude = + std::ldexp(static_cast(mantissa | 0x400u), exponent - 25); + } + uint32_t bits; + std::memcpy(&bits, &magnitude, 4); + bits |= sign; + float out; + std::memcpy(&out, &bits, 4); + return out; +} + +float bfloat16_to_float(uint16_t b) { + const uint32_t bits = static_cast(b) << 16; + float out; + std::memcpy(&out, &bits, 4); + return out; +} +} // namespace + +std::pair prepare_text(const std::string& text) { + const auto words = split_words(text); + std::string out = join(words); + if (out.empty()) { + return {out, 3}; + } + const int frames_after_eos = words.size() <= 4 ? 3 : 1; + if (out[0] >= 'a' && out[0] <= 'z') { + out[0] = static_cast(out[0] - 'a' + 'A'); + } + if (is_ascii_alnum(out.back())) { + out += '.'; + } + return {out, frames_after_eos}; +} + +std::vector split_sentences( + const std::string& text, + const std::function& count_tokens, + size_t max_tokens) { + std::vector chunks; + std::string current; + for (const auto& sentence : split_at_sentence_ends(join(split_words(text)))) { + std::vector pieces{sentence}; + if (count_tokens(sentence) > max_tokens) { + pieces = split_long_sentence(sentence, count_tokens, max_tokens); + } + for (const auto& piece : pieces) { + std::string candidate = join(current, piece); + if (!current.empty() && count_tokens(candidate) > max_tokens) { + chunks.push_back(current); + candidate = piece; + } + current = std::move(candidate); + } + } + if (!current.empty()) { + chunks.push_back(current); + } + return chunks; +} + +double NoiseGenerator::uniform() { + state_ += 0x9e3779b97f4a7c15ull; + uint64_t z = state_; + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebull; + z ^= z >> 31; + return static_cast(z >> 11) * 0x1.0p-53; +} + +void NoiseGenerator::fill(float* out, size_t count, double stddev) { + constexpr double kTwoPi = 6.283185307179586; + for (size_t i = 0; i < count; i += 2) { + const double radius = std::sqrt(-2.0 * std::log(1.0 - uniform())); + const double angle = kTwoPi * uniform(); + out[i] = static_cast(radius * std::cos(angle) * stddev); + if (i + 1 < count) { + out[i + 1] = static_cast(radius * std::sin(angle) * stddev); + } + } +} + +Runner::Runner( + const std::string& model_path, + const std::string& tokenizer_path, + const std::string& data_path) + : model_(model_path, data_path, Module::LoadMode::Mmap) { + if (model_.load() != Error::Ok) { + throw std::runtime_error("Could not load " + model_path); + } + if (tokenizer_.load(tokenizer_path) != tokenizers::Error::Ok) { + throw std::runtime_error("Could not load tokenizer " + tokenizer_path); + } + const auto integer = [this](const char* name) { + return execute(model_, name).at(0).toInt(); + }; + const auto real = [this](const char* name) { + return execute(model_, name).at(0).toDouble(); + }; + if (integer("format_version") != kFormatVersion) { + throw std::runtime_error("Unsupported Phonon program; re-export it"); + } + sample_rate_ = integer("sample_rate"); + frame_samples_ = integer("frame_samples"); + dim_ = integer("dim"); + latent_dim_ = integer("latent_dim"); + max_context_ = integer("max_context"); + max_prefill_ = integer("max_prefill"); + frame_rate_ = real("frame_rate"); + eos_threshold_ = real("eos_threshold"); + temperature_ = real("temperature"); + bos_embeds_ = copy_floats(execute(model_, "bos_embeds").at(0)); + if (static_cast(bos_embeds_.size()) != dim_) { + throw std::runtime_error("bos_embeds does not match dim"); + } + noise_.assign(latent_dim_, 0.0f); + for (const char* method : + {"text_embed", "voice_embed", "flow_lm", "mimi_decode"}) { + if (model_.load_method(method) != Error::Ok) { + throw std::runtime_error(std::string("Could not load method ") + method); + } + } + auto meta = model_.method_meta("mimi_decode"); + if (!meta.ok()) { + throw std::runtime_error("Could not read mimi_decode metadata"); + } + for (size_t i = 1; i < meta->num_inputs(); ++i) { + auto info = meta->input_tensor_meta(i); + if (!info.ok() || info->scalar_type() != ScalarType::Float) { + throw std::runtime_error("mimi_decode state must be float32 tensors"); + } + const auto sizes = info->sizes(); + mimi_state_sizes_.emplace_back(sizes.begin(), sizes.end()); + } +} + +std::vector Runner::tokenize(const std::string& text) const { + auto tokens = tokenizer_.encode(text, /*bos=*/0, /*eos=*/0); + if (!tokens.ok()) { + throw std::runtime_error("Could not tokenize: " + text); + } + return std::move(tokens.get()); +} + +Voice Runner::load_voice(const std::string& path) const { + std::ifstream file(path, std::ios::binary); + uint64_t header_size = 0; + if (!file.read(reinterpret_cast(&header_size), 8) || + header_size > (1u << 24)) { + throw std::runtime_error("Not a safetensors file: " + path); + } + std::string header(header_size, '\0'); + file.read(header.data(), static_cast(header_size)); + const auto json = nlohmann::json::parse(header); + if (!json.contains("speaker_wavs")) { + throw std::runtime_error(path + ": expected a tensor named speaker_wavs"); + } + const auto& entry = json.at("speaker_wavs"); + const auto dtype = entry.at("dtype").get(); + const auto shape = entry.at("shape").get>(); + const auto offsets = entry.at("data_offsets").get>(); + if (shape.size() != 3 || shape[0] != 1 || shape[1] != dim_ || shape[2] <= 0 || + shape[2] > max_prefill_) { + throw std::runtime_error( + path + ": speaker_wavs must be [1, " + std::to_string(dim_) + + ", frames] with at most " + std::to_string(max_prefill_) + " frames"); + } + const size_t element_size = dtype == "F32" ? 4 : 2; + if (dtype != "F32" && dtype != "F16" && dtype != "BF16") { + throw std::runtime_error(path + ": unsupported dtype " + dtype); + } + const int64_t frames = shape[2]; + const size_t count = static_cast(dim_ * frames); + if (offsets.size() != 2 || offsets[1] - offsets[0] != count * element_size) { + throw std::runtime_error(path + ": inconsistent speaker_wavs size"); + } + std::vector raw(count * element_size); + file.seekg(static_cast(8 + header_size + offsets[0])); + if (!file.read( + reinterpret_cast(raw.data()), + static_cast(raw.size()))) { + throw std::runtime_error(path + ": truncated file"); + } + const auto element = [&](size_t k) -> float { + const unsigned char* p = raw.data() + k * element_size; + if (dtype == "F32") { + float value; + std::memcpy(&value, p, 4); + return value; + } + return dtype == "F16" ? half_to_float(read_u16_le(p)) + : bfloat16_to_float(read_u16_le(p)); + }; + Voice voice; + voice.frames = frames; + voice.latents.resize(count); + for (int64_t d = 0; d < dim_; ++d) { + for (int64_t t = 0; t < frames; ++t) { + voice.latents[t * dim_ + d] = element(d * frames + t); + } + } + return voice; +} + +Runner::FlowOutput Runner::flow(const std::vector& embeds, int64_t pos) { + return flow(embeds.data(), static_cast(embeds.size()) / dim_, pos); +} + +Runner::FlowOutput Runner::flow(const float* embeds, int64_t seq, int64_t pos) { + int64_t input_pos = pos; + auto embeds_tensor = from_blob( + const_cast(embeds), + {1, static_cast(seq), static_cast(dim_)}); + auto pos_tensor = from_blob(&input_pos, {1}, ScalarType::Long); + auto noise_tensor = + from_blob(noise_.data(), {1, static_cast(latent_dim_)}); + auto out = + execute(model_, "flow_lm", {embeds_tensor, pos_tensor, noise_tensor}); + return { + copy_floats(out.at(0)), float_data(out.at(1))[0], copy_floats(out.at(2))}; +} + +Synthesis Runner::synthesize( + const std::string& text, + const Voice& voice, + double temperature, + uint64_t seed, + const std::function& on_audio) { + if (voice.frames <= 0 || voice.frames > max_prefill_ || + static_cast(voice.latents.size()) != voice.frames * dim_) { + throw std::invalid_argument("Invalid voice prompt"); + } + NoiseGenerator noise(seed); + const double stddev = std::sqrt(temperature < 0 ? temperature_ : temperature); + + auto voice_tensor = from_blob( + const_cast(voice.latents.data()), + {1, static_cast(voice.frames), static_cast(dim_)}); + const auto voice_embeds = + copy_floats(execute(model_, "voice_embed", {voice_tensor}).at(0)); + std::fill(noise_.begin(), noise_.end(), 0.0f); + flow(voice_embeds, 0); + const int64_t voice_len = voice.frames; + + std::vector> state; + for (const auto& sizes : mimi_state_sizes_) { + int64_t numel = 1; + for (auto s : sizes) { + numel *= s; + } + state.emplace_back(numel, 0.0f); + } + + Synthesis result; + const auto count_tokens = [this](const std::string& s) { + return tokenize(s).size(); + }; + for (const auto& chunk_text : split_sentences( + prepare_text(text).first, count_tokens, kMaxChunkTokens)) { + auto [prepared, frames_after_eos] = prepare_text(chunk_text); + Chunk chunk{prepared, tokenize(prepared), 0}; + const auto num_tokens = static_cast(chunk.tokens.size()); + if (num_tokens == 0 || num_tokens > max_prefill_) { + throw std::runtime_error( + "Chunk has an unsupported token count: " + prepared); + } + std::vector ids(chunk.tokens.begin(), chunk.tokens.end()); + auto tokens_tensor = from_blob( + ids.data(), {1, static_cast(num_tokens)}, ScalarType::Long); + const auto text_embeds = + copy_floats(execute(model_, "text_embed", {tokens_tensor}).at(0)); + std::fill(noise_.begin(), noise_.end(), 0.0f); + int64_t pos = voice_len; + flow(text_embeds, pos); + pos += num_tokens; + + const int64_t budget = std::min( + frame_budget(chunk.tokens.size(), frame_rate_), max_context_ - pos); + std::vector embeds = bos_embeds_; + int countdown = -1; + for (int64_t k = 0; k < budget; ++k) { + noise.fill(noise_.data(), noise_.size(), stddev); + auto step = flow(embeds, pos); + embeds = std::move(step.next_embeds); + ++pos; + ++chunk.frames; + + // The tensors must outlive the EValues that point to them. + std::vector tensors; + tensors.reserve(state.size() + 1); + tensors.push_back(from_blob( + step.latent.data(), {1, 1, static_cast(latent_dim_)})); + for (size_t i = 0; i < state.size(); ++i) { + tensors.push_back(from_blob(state[i].data(), mimi_state_sizes_[i])); + } + auto out = execute( + model_, + "mimi_decode", + std::vector(tensors.begin(), tensors.end())); + const float* pcm = float_data(out.at(0)); + const size_t pcm_size = out.at(0).toTensor().numel(); + result.pcm.insert(result.pcm.end(), pcm, pcm + pcm_size); + if (on_audio) { + on_audio(pcm, pcm_size); + } + for (size_t i = 0; i < state.size(); ++i) { + std::memcpy( + state[i].data(), float_data(out.at(i + 1)), state[i].size() * 4); + } + + if (step.eos_logit > eos_threshold_ && countdown < 0) { + countdown = frames_after_eos; + } + if (countdown >= 0) { + if (countdown == 0) { + break; + } + --countdown; + } + } + result.frames += chunk.frames; + result.chunks.push_back(std::move(chunk)); + } + return result; +} + +} // namespace phonon diff --git a/examples/models/phonon/runner.h b/examples/models/phonon/runner.h new file mode 100644 index 00000000000..646696b91be --- /dev/null +++ b/examples/models/phonon/runner.h @@ -0,0 +1,111 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace phonon { + +// Collapses whitespace, capitalizes the first letter and ends the text with a +// period if it ends with a letter or digit. Returns the text and the number of +// frames to generate after the end-of-speech logit fires (3 for texts of up to +// 4 words, else 1). Same rules as run_phonon.py. +std::pair prepare_text(const std::string& text); + +// Packs whole sentences into chunks of at most max_tokens tokens; a longer +// sentence is split at word boundaries. Same rules as run_phonon.py. +std::vector split_sentences( + const std::string& text, + const std::function& count_tokens, + size_t max_tokens); + +// Gaussian noise for the flow head, identical to NoiseGenerator in +// run_phonon.py: splitmix64 uniforms and the Box-Muller transform. +class NoiseGenerator { + public: + explicit NoiseGenerator(uint64_t seed) : state_(seed) {} + void fill(float* out, size_t count, double stddev); + + private: + double uniform(); + uint64_t state_; +}; + +// Speaker latents, time-major [frames, dim]. +struct Voice { + std::vector latents; + int64_t frames = 0; +}; + +struct Chunk { + std::string text; + std::vector tokens; + int64_t frames = 0; +}; + +struct Synthesis { + std::vector pcm; + std::vector chunks; + int64_t frames = 0; +}; + +class Runner { + public: + // data_path is the external data file of the program, if any + // (aoti_cuda_blob.ptd for CUDA). + Runner( + const std::string& model_path, + const std::string& tokenizer_path, + const std::string& data_path = ""); + + // Reads a Phonon voice .safetensors file (speaker_wavs [1, dim, frames]). + Voice load_voice(const std::string& path) const; + + // temperature < 0 uses the model's default. on_audio, if set, receives the + // PCM of every frame as soon as it is decoded. + Synthesis synthesize( + const std::string& text, + const Voice& voice, + double temperature = -1, + uint64_t seed = 0, + const std::function& on_audio = nullptr); + + std::vector tokenize(const std::string& text) const; + + int64_t sample_rate() const { + return sample_rate_; + } + + private: + struct FlowOutput { + std::vector latent; + float eos_logit; + std::vector next_embeds; + }; + FlowOutput flow(const std::vector& embeds, int64_t pos); + FlowOutput flow(const float* embeds, int64_t seq, int64_t pos); + + executorch::extension::Module model_; + tokenizers::SPTokenizer tokenizer_; + int64_t sample_rate_, frame_samples_, dim_, latent_dim_, max_context_, + max_prefill_; + double frame_rate_, eos_threshold_, temperature_; + std::vector bos_embeds_, noise_; + std::vector> mimi_state_sizes_; +}; + +} // namespace phonon diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py new file mode 100644 index 00000000000..69a2c228a54 --- /dev/null +++ b/examples/models/phonon/source_transformations.py @@ -0,0 +1,231 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Function-preserving module swaps for backends that lack an op. + +The model in model.py uses plain PyTorch ops. Each swap below computes the +same function from more widely supported ops; export_phonon.py selects them +per backend: + + decompose_elu MLX and Vulkan have no aten.elu kernel; without it + the Mimi decoder splits at every ELU. + conv_transpose1d_as_2d Vulkan runs transposed convolution on 4-D inputs only. + depthwise_conv_transpose_as_mul + Vulkan computes the Mimi upsampler (depthwise, + kernel = 2 * stride) incorrectly; this runs it as a + multiply and an overlap-add. + use_webgpu_compat_ops WebGPU runtimes built before these runtime fixes + (off by default; --webgpu-compat-ops): + * native_layer_norm resize hook (the flow LM's + dynamic sequence length), + * reduce ops reading a Bool keepdim argument, + * aten.elu and aten.silu kernels, + * 1-D transposed convolution. + compute_weights_in --dtype bf16 (CUDA and MLX; FP32 is the default): + linear, convolution and embedding weights and their + matmuls in bf16; norms, attention, RoPE and the + streaming state stay in float32. Not + function-preserving. + +Each function swaps modules in place in the programs from build_programs() and +returns how many it replaced. +""" + +import torch +import torch.nn.functional as F + +from executorch.examples.models.phonon.model import FlowLMStep, StreamingConvTranspose1d +from torch import nn + + +class LayerNormFromRMSNorm(nn.Module): + """LayerNorm whose mean is a linear with weight 1/D (exact for D a power + of two) followed by the RMS-norm pattern, which follows a resized sequence + length without a native_layer_norm or mean.dim resize hook.""" + + def __init__(self, norm: nn.LayerNorm): + super().__init__() + (dim,) = norm.normalized_shape + self.weight, self.bias, self.eps = norm.weight, norm.bias, norm.eps + self.register_buffer("mean_weight", torch.full((1, dim), 1.0 / dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + xc = x - F.linear(x, self.mean_weight) + rms = torch.rsqrt((xc * xc).mean(-1, keepdim=True) + self.eps) + return xc * rms * self.weight + self.bias + + +class ELUFromExp(nn.Module): + """ELU(alpha=1) = relu(x) + exp(min(x, 0)) - 1.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x) + torch.exp(torch.clamp(x, max=0.0)) - 1.0 + + +class SiLUFromSigmoid(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(x) + + +class ConvTranspose1dAs2d(nn.Module): + """ConvTranspose1d as a conv_transpose2d over a height-1 image. The 4-D + weight is stored as a parameter: Vulkan prepacks only constant weights.""" + + def __init__(self, convtr: nn.ConvTranspose1d): + super().__init__() + self.weight = nn.Parameter( + convtr.weight.detach().unsqueeze(2), requires_grad=False + ) + self.bias = convtr.bias + self.stride, self.groups = convtr.stride[0], convtr.groups + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = F.conv_transpose2d( + x.unsqueeze(2), + self.weight, + self.bias, + stride=(1, self.stride), + groups=self.groups, + ) + return y.squeeze(2) + + +class ComputeIn(nn.Module): + """Runs `module` with its weights in `dtype` on a float32 input and + returns float32. Other attributes (bias, kernel_size, ...) are read from + the wrapped module.""" + + def __init__(self, module: nn.Module, dtype: torch.dtype): + super().__init__() + self.inner = module.to(dtype) + self.dtype = dtype + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("inner"), name) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.is_floating_point(): + return self.inner(x.to(self.dtype)).to(x.dtype) + return self.inner(x).float() + + +class DepthwiseConvTranspose1dAsMul(nn.Module): + """Depthwise ConvTranspose1d with kernel_size == 2 * stride, without bias: + each input step scales the kernel, and the halves of neighboring steps + overlap.""" + + def __init__(self, convtr: nn.ConvTranspose1d): + super().__init__() + (self.stride,) = convtr.stride + channels, _, kernel = convtr.weight.shape + self.weight = nn.Parameter( + convtr.weight.detach().reshape(1, channels, 1, kernel), requires_grad=False + ) + self.bias = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, c, n = x.shape + s = self.stride + y = x.unsqueeze(3) * self.weight # [b, c, n, 2s] + head = y[..., :s].reshape(b, c, n * s) + tail = y[..., s:].reshape(b, c, n * s) + zeros = torch.zeros(b, c, s, dtype=x.dtype) + return torch.cat([head, zeros], 2) + torch.cat([zeros, tail], 2) + + +def _swap(root: nn.Module, cls: type, make) -> int: + count = 0 + for parent in list(root.modules()): + for name, child in list(parent.named_children()): + if isinstance(child, cls): + setattr(parent, name, make(child)) + count += 1 + return count + + +def decompose_elu(programs: dict[str, nn.Module]) -> int: + return sum(_swap(p, nn.ELU, lambda m: ELUFromExp()) for p in programs.values()) + + +def decompose_silu(programs: dict[str, nn.Module]) -> int: + return sum( + _swap(p, nn.SiLU, lambda m: SiLUFromSigmoid()) for p in programs.values() + ) + + +def conv_transpose1d_as_2d(programs: dict[str, nn.Module]) -> int: + count = 0 + for program in programs.values(): + for module in program.modules(): + if isinstance(module, StreamingConvTranspose1d) and isinstance( + module.convtr, nn.ConvTranspose1d + ): + module.convtr = ConvTranspose1dAs2d(module.convtr) + count += 1 + return count + + +def depthwise_conv_transpose_as_mul(programs: dict[str, nn.Module]) -> int: + count = 0 + for program in programs.values(): + for module in program.modules(): + convtr = ( + getattr(module, "convtr", None) + if isinstance(module, StreamingConvTranspose1d) + else None + ) + if ( + isinstance(convtr, nn.ConvTranspose1d) + and convtr.bias is None + and convtr.groups == convtr.in_channels == convtr.out_channels + and convtr.kernel_size[0] == 2 * convtr.stride[0] + ): + module.convtr = DepthwiseConvTranspose1dAsMul(convtr) + count += 1 + return count + + +def decompose_dynamic_layer_norm(programs: dict[str, nn.Module]) -> int: + """Only the flow LM norms that see the dynamic sequence length.""" + count = 0 + for program in programs.values(): + if isinstance(program, FlowLMStep): + for layer in program.layers: + layer.norm1 = LayerNormFromRMSNorm(layer.norm1) + layer.norm2 = LayerNormFromRMSNorm(layer.norm2) + count += 2 + return count + + +WEIGHT_MODULES = ( + nn.Linear, + nn.Conv1d, + nn.ConvTranspose1d, + nn.Embedding, + ConvTranspose1dAs2d, +) + + +def compute_weights_in(programs: dict[str, nn.Module], dtype: torch.dtype) -> int: + return sum( + _swap(p, WEIGHT_MODULES, lambda m: ComputeIn(m, dtype)) + for p in programs.values() + ) + + +def use_webgpu_compat_ops(programs: dict[str, nn.Module]) -> int: + return sum( + swap(programs) + for swap in ( + decompose_dynamic_layer_norm, + decompose_elu, + decompose_silu, + conv_transpose1d_as_2d, + ) + ) diff --git a/examples/models/phonon/tests/__init__.py b/examples/models/phonon/tests/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/examples/models/phonon/tests/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/examples/models/phonon/tests/reference.py b/examples/models/phonon/tests/reference.py new file mode 100644 index 00000000000..e13b79aa8a7 --- /dev/null +++ b/examples/models/phonon/tests/reference.py @@ -0,0 +1,162 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Straightforward eager Phonon for tests: interleaved RoPE on the original +in_proj layout, a growing KV cache, and Python-side streaming state, written +to follow the official runtime operation by operation. The export programs in +model.py must compute the same function.""" + +import math + +import torch +import torch.nn.functional as F + +from executorch.examples.models.phonon.model import Phonon, rope_inv_freq + + +def rope_interleaved(x, pos, max_period): + """x [B, T, H, hd], absolute positions pos [T]; pairs (2i, 2i+1) rotate.""" + inv = rope_inv_freq(x.shape[-1], max_period) + ang = pos.to(torch.float32)[:, None] * inv[None, :] + c, s = torch.cos(ang)[None, :, None, :], torch.sin(ang)[None, :, None, :] + xr, xi = x[..., 0::2], x[..., 1::2] + return torch.stack([xr * c - xi * s, xr * s + xi * c], dim=-1).flatten(-2) + + +class FlowLMReference: + def __init__(self, model: Phonon): + assert model.rope_layout == "interleaved" + self.m, self.cfg = model.flow_lm, model.cfg + with torch.no_grad(): + self.cond = self.m.summed_conditions() + self.t_cond = self.m.flow_net.lsd_time_condition() + self.k = [None] * self.cfg.num_layers + self.v = [None] * self.cfg.num_layers + self.pos = 0 + + @torch.no_grad() + def backbone(self, x): + cfg = self.cfg + _, t, d = x.shape + h, hd = cfg.num_heads, cfg.head_dim + pos = torch.arange(self.pos, self.pos + t) + for i, layer in enumerate(self.m.transformer.layers): + qkv = layer.self_attn.in_proj(layer.norm1(x)).view(1, t, 3, h, hd) + q = rope_interleaved(qkv[:, :, 0], pos, cfg.max_period) + k = rope_interleaved(qkv[:, :, 1], pos, cfg.max_period) + v = qkv[:, :, 2] + self.k[i] = k if self.k[i] is None else torch.cat([self.k[i], k], 1) + self.v[i] = v if self.v[i] is None else torch.cat([self.v[i], v], 1) + kk, vv, qq = (z.transpose(1, 2) for z in (self.k[i], self.v[i], q)) + n = kk.shape[2] + qpos = torch.arange(n - t, n)[:, None] + mask = torch.where(torch.arange(n)[None, :] <= qpos, 0.0, float("-inf")) + att = torch.softmax(qq @ kk.transpose(-1, -2) / math.sqrt(hd) + mask, -1) + o = (att @ vv).transpose(1, 2).reshape(1, t, d) + x = x + layer.self_attn.out_proj(o) + x = x + layer.linear2(F.gelu(layer.linear1(layer.norm2(x)))) + self.pos += t + return x + + def prompt(self, embeds): + self.backbone(embeds) + + @torch.no_grad() + def step(self, prev_latent, noise): + """prev_latent None (BOS) or [1, 1, ldim]; returns latent, eos_logit.""" + m = self.m + seq = m.bos_emb.view(1, 1, -1) if prev_latent is None else prev_latent + x = m.input_linear(seq) + self.cond + h = m.out_norm(self.backbone(x))[:, -1] + latent = noise + m.flow_net(h, noise, self.t_cond) + return latent.view(1, 1, -1), m.out_eos(h) + + +def _conv_stream(conv, x, prev): + xx = x if prev is None else torch.cat([prev, x], 2) + y = F.conv1d(xx, conv.weight, conv.bias) + if prev is None: + return y, None + return y, xx[:, :, xx.shape[2] - prev.shape[2] :] + + +def _convtr_stream(convtr, x, partial): + y = F.conv_transpose1d( + x, convtr.weight, convtr.bias, stride=convtr.stride, groups=convtr.groups + ) + pt = partial.shape[2] + y = torch.cat([y[:, :, :pt] + partial, y[:, :, pt:]], 2) + new_partial = y[:, :, y.shape[2] - pt :] + if convtr.bias is not None: + new_partial = new_partial - convtr.bias.view(1, -1, 1) + return y[:, :, : y.shape[2] - pt], new_partial + + +class MimiReference: + """Streaming Mimi decoder; decode() takes any number of frames.""" + + def __init__(self, model: Phonon): + assert model.rope_layout == "interleaved" + cfg, m = model.cfg, model.mimi + self.cfg, self.m, self.lm = cfg, m, model.flow_lm + z = lambda *s: torch.zeros(1, *s) # noqa: E731 + self.up = z(cfg.mimi_dim, m.upsample.convtr.state_len) + self.kv = [[None, None] for _ in range(cfg.mimi_layers)] + self.offset = 0 + dec = m.decoder.model + self.prev = {0: z(cfg.mimi_dim, dec[0].state_len)} + self.partial = {} + for j in range(len(cfg.ratios)): + up, res = dec[2 + 3 * j], dec[3 + 3 * j] + self.partial[j] = z(up.convtr.out_channels, up.state_len) + self.prev[("res", j)] = z(up.convtr.out_channels, res.block[1].state_len) + self.prev["final"] = z(cfg.n_filters, dec[-1].state_len) + + @torch.no_grad() + def transformer(self, x): + cfg = self.cfg + t, d, h, hd = x.shape[1], cfg.mimi_dim, cfg.mimi_heads, cfg.mimi_head_dim + pos = torch.arange(self.offset, self.offset + t) + for i, layer in enumerate(self.m.decoder_transformer.transformer.layers): + qkv = layer.self_attn.in_proj(layer.norm1(x)).view(1, t, 3, h, hd) + q = rope_interleaved(qkv[:, :, 0], pos, cfg.mimi_max_period).transpose(1, 2) + k = rope_interleaved(qkv[:, :, 1], pos, cfg.mimi_max_period).transpose(1, 2) + v = qkv[:, :, 2].transpose(1, 2) + pk, pv = self.kv[i] + kk = k if pk is None else torch.cat([pk, k], 2) + vv = v if pv is None else torch.cat([pv, v], 2) + qi = torch.arange(t)[:, None] + kk.shape[2] - t + kj = torch.arange(kk.shape[2])[None, :] + ok = (kj <= qi) & (kj >= qi - cfg.mimi_context) + mask = torch.where(ok, 0.0, float("-inf")) + att = torch.softmax(q @ kk.transpose(-1, -2) / math.sqrt(hd) + mask, -1) + o = (att @ vv).transpose(1, 2).reshape(1, t, d) + self.kv[i] = [kk[:, :, -cfg.mimi_context :], vv[:, :, -cfg.mimi_context :]] + x = x + layer.layer_scale_1.scale * layer.self_attn.out_proj(o) + h2 = layer.linear2(F.gelu(layer.linear1(layer.norm2(x)))) + x = x + layer.layer_scale_2.scale * h2 + self.offset += t + return x + + @torch.no_grad() + def decode(self, latent): + """latent [1, T, ldim] -> pcm [1, 1, T * frame_samples].""" + m, dec = self.m, self.m.decoder.model + x = (latent * self.lm.emb_std + self.lm.emb_mean).transpose(1, 2) + x = F.conv1d(x, m.quantizer.output_proj.weight) + x, self.up = _convtr_stream(m.upsample.convtr.convtr, x, self.up) + x = self.transformer(x.transpose(1, 2)).transpose(1, 2) + x, self.prev[0] = _conv_stream(dec[0].conv, x, self.prev[0]) + for j in range(len(self.cfg.ratios)): + up, res = dec[2 + 3 * j], dec[3 + 3 * j] + x, self.partial[j] = _convtr_stream(up.convtr, F.elu(x), self.partial[j]) + h, self.prev[("res", j)] = _conv_stream( + res.block[1].conv, F.elu(x), self.prev[("res", j)] + ) + h, _ = _conv_stream(res.block[3].conv, F.elu(h), None) + x = x + h + x, self.prev["final"] = _conv_stream(dec[-1].conv, F.elu(x), self.prev["final"]) + return x diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py new file mode 100644 index 00000000000..5553d4ee4ed --- /dev/null +++ b/examples/models/phonon/tests/test_phonon.py @@ -0,0 +1,409 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Fast checks on a tiny random Phonon: the export programs against the eager +reference, the backend module swaps, checkpoint loading, the I/O contract of +exported programs, and the host-side text handling and noise. + + python -m pytest examples/models/phonon/tests -q +""" + +import copy +import json +import tempfile +import unittest + +import torch + +from executorch.examples.models.phonon.export_phonon import ( + lower_phonon, + quantize_scheme, + source_transforms_for, +) +from executorch.examples.models.phonon.model import ( + build_programs, + example_inputs, + io_contract, + mimi_zero_state, + Phonon, + PhononConfig, + program_metadata, +) +from executorch.examples.models.phonon.run_phonon import ( + EagerPrograms, + frame_budget, + NoiseGenerator, + PhononRunner, + prepare_text, + split_sentences, +) +from executorch.examples.models.phonon.source_transformations import ( + compute_weights_in, + use_webgpu_compat_ops, +) +from executorch.examples.models.phonon.tests.reference import ( + FlowLMReference, + MimiReference, +) + +TINY = PhononConfig( + dim=64, + num_heads=2, + num_layers=2, + ffn_dim=128, + text_bins=50, + latent_dim=8, + flow_depth=2, + duration_dim=16, + sample_rate=1600, + frame_rate=50.0, + mimi_dim=32, + n_filters=8, + ratios=(2, 2), + mimi_heads=2, + mimi_layers=2, + mimi_ffn_dim=64, + mimi_context=12, + max_context=64, + max_prefill=16, +) + + +def random_model(cfg: PhononConfig = TINY, seed: int = 0) -> Phonon: + """A Phonon with every parameter random (interleaved RoPE layout).""" + torch.manual_seed(seed) + model = Phonon(cfg).eval().requires_grad_(False) + for p in model.parameters(): + p.copy_(torch.randn_like(p) * 0.2) + return model + + +def official_config() -> dict: + """The fields of gradium/phonon-31m-en's config.json that are parsed.""" + transformer = { + "transformer_num_heads": 8, + "transformer_num_layers": 2, + "transformer_context": 250, + "transformer_max_period": 10000.0, + "transformer_dim_feedforward": 2048, + } + return { + "languages": ["en"], + "flow_lm": { + "d_model": 512, + "num_heads": 8, + "num_layers": 6, + "dim_feedforward": 2048, + "max_period": 10000.0, + "n_bins": 4000, + "lut_dim": 512, + "flow_dim": 512, + "flow_depth": 6, + "ldim": 32, + "time_rms_norm": "rms", + "gelu": "erf", + }, + "mimi": { + "channels": 1, + "sample_rate": 24000, + "frame_rate": 12.5, + "dimension": 512, + "n_filters": 64, + "n_residual_layers": 1, + "ratios": [6, 5, 4], + "kernel_size": 7, + "last_kernel_size": 3, + "residual_kernel_size": 3, + "compress": 2, + "gelu": "erf", + **transformer, + }, + "lsd_decode_steps": 1, + "eos_threshold": 4.0, + "conditioners": [ + { + "type": "lut", + "lut": {"n_bins": 31, "dim": 16, "possible_values": ["0", "1", "2"]}, + "name": "num_speakers", + }, + { + "type": "continuous", + "continuous": { + "scale_factor": 1000.0, + "dim": 128, + "max_period": 10000.0, + }, + "name": "duration_delta", + }, + ], + "fuser": {"sum": ["num_speakers", "duration_delta"]}, + } + + +def max_diff(a: torch.Tensor, b: torch.Tensor) -> float: + return (a - b).abs().max().item() + + +class TestConfig(unittest.TestCase): + def test_official_config_matches_defaults(self): + cfg = PhononConfig.from_dict(official_config()) + self.assertEqual(cfg, PhononConfig()) + self.assertEqual(cfg.frame_samples, 1920) + self.assertEqual(cfg.upsample_stride, 16) + + def test_rejects_unsupported_variant(self): + data = official_config() + data["lsd_decode_steps"] = 4 + with self.assertRaises(ValueError): + PhononConfig.from_dict(data) + + +class TestPrograms(unittest.TestCase): + def setUp(self): + self.ref_model = random_model() + self.model = copy.deepcopy(self.ref_model).to_half_split_rope() + + def _check_flow_lm(self, attention: str): + cfg = TINY + flow = build_programs(self.model, attention)["flow_lm"] + ref = FlowLMReference(self.ref_model) + g = torch.Generator().manual_seed(1) + voice = torch.randn(1, 6, cfg.dim, generator=g) + text = torch.randn(1, 4, cfg.dim, generator=g) + zeros = torch.zeros(1, cfg.latent_dim) + with torch.no_grad(): + flow(voice, torch.tensor([0]), zeros) + flow(text, torch.tensor([6]), zeros) + ref.prompt(voice) + ref.prompt(text) + embeds, prev, pos = self.model.flow_lm.bos_embeds(), None, 10 + for _ in range(6): + noise = torch.randn(1, cfg.latent_dim, generator=g) + latent, eos, embeds = flow(embeds, torch.tensor([pos]), noise) + ref_latent, ref_eos = ref.step(prev, noise) + self.assertLess(max_diff(latent, ref_latent.view(1, -1)), 1e-4) + self.assertLess(max_diff(eos, ref_eos), 1e-4) + prev, pos = ref_latent, pos + 1 + + def test_flow_lm_custom_sdpa_matches_reference(self): + self._check_flow_lm("custom_sdpa") + + def test_flow_lm_sdpa_matches_reference(self): + self._check_flow_lm("sdpa") + + def test_mimi_matches_reference(self): + cfg = TINY + mimi = build_programs(self.model)["mimi_decode"] + ref = MimiReference(self.ref_model) + g = torch.Generator().manual_seed(2) + state = mimi_zero_state(cfg) + # Enough frames to slide the transformer window several times. + for _ in range(6): + latent = torch.randn(1, 1, cfg.latent_dim, generator=g) + with torch.no_grad(): + pcm, *state = mimi(latent, *state) + self.assertEqual(tuple(pcm.shape), (1, 1, cfg.frame_samples)) + self.assertLess(max_diff(pcm, ref.decode(latent)), 1e-4) + + def test_weight_names_are_unique_across_programs(self): + # The CUDA backend shares a weight between methods when the names match. + programs = build_programs(self.model, "sdpa") + compute_weights_in(programs, torch.bfloat16) + owner = {} + for name, program in programs.items(): + for fqn, tensor in [*program.named_parameters(), *program.named_buffers()]: + if fqn in owner and owner[fqn][1] is not tensor: + self.fail(f"{fqn} is in {owner[fqn][0]} and {name}") + owner[fqn] = (name, tensor) + + def test_webgpu_compat_ops_preserve_outputs(self): + inputs = example_inputs(TINY) + programs = build_programs(self.model) + with torch.no_grad(): + before = {n: programs[n](*inputs[n]) for n in ("flow_lm", "mimi_decode")} + # norm1/norm2 x 2 layers; 2 x 2 residual + 3 decoder ELUs; 2 time + # embedder + 2 x 2 res block + 1 final SiLUs; 3 transposed convs. + self.assertEqual(use_webgpu_compat_ops(programs), 4 + 7 + 7 + 3) + for p in programs.values(): + if hasattr(p, "reset"): + p.reset() + with torch.no_grad(): + for name, outputs in before.items(): + for a, b in zip(outputs, programs[name](*inputs[name])): + self.assertLess(max_diff(a, b), 1e-4) + + def test_vulkan_swaps_preserve_outputs(self): + inputs = example_inputs(TINY) + programs = build_programs(self.model) + with torch.no_grad(): + before = programs["mimi_decode"](*inputs["mimi_decode"]) + # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. + counts = [t(programs) for t in source_transforms_for("vulkan", False)] + self.assertEqual(counts, [7, 1, 2]) + with torch.no_grad(): + after = programs["mimi_decode"](*inputs["mimi_decode"]) + for a, b in zip(before, after): + self.assertLess(max_diff(a, b), 1e-4) + + def test_compute_weights_in_keeps_float32_io(self): + inputs = example_inputs(TINY) + programs = build_programs(self.model) + with torch.no_grad(): + before = programs["mimi_decode"](*inputs["mimi_decode"]) + self.assertGreater(compute_weights_in(programs, torch.bfloat16), 0) + with torch.no_grad(): + after = programs["mimi_decode"](*inputs["mimi_decode"]) + for a, b in zip(before, after): + self.assertEqual(b.dtype, torch.float32) + self.assertLess(max_diff(a, b), 0.05 * (a.abs().max().item() + 1e-6)) + + +class TestCheckpoint(unittest.TestCase): + def test_strict_load_drops_only_training_tensors(self): + from safetensors.torch import save_file + + model = random_model() + state = {k: v.to(torch.bfloat16) for k, v in model.state_dict().items()} + state["flow_lm.flow.w_s_t.0.weight"] = torch.zeros(32, 2) + with tempfile.TemporaryDirectory() as d: + save_file(state, f"{d}/model.safetensors") + loaded = Phonon.from_checkpoint(f"{d}/model.safetensors", TINY) + self.assertEqual(loaded.rope_layout, "half_split") + del state["flow_lm.out_eos.bias"] + save_file(state, f"{d}/model.safetensors") + with self.assertRaises(RuntimeError): + Phonon.from_checkpoint(f"{d}/model.safetensors", TINY) + + +class TestExport(unittest.TestCase): + def _lower(self, backend: str): + model = random_model().to_half_split_rope() + program, delegation, metadata = lower_phonon(model, backend) + return model, program, delegation, metadata + + def test_portable_program_matches_contract_and_eager(self): + model, program, _, metadata = self._lower("portable") + try: + import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, + ) + except ImportError: + self.skipTest("ExecuTorch pybindings are not available") + module = _load_for_executorch_from_buffer(program.buffer) + contract = io_contract(TINY) + for method, spec in contract.items(): + meta = module.method_meta(method) + self.assertEqual(meta.num_inputs(), len(spec["inputs"])) + self.assertEqual(meta.num_outputs(), len(spec["outputs"])) + for i, tensor in enumerate(spec["inputs"]): + sizes = list(meta.input_tensor_meta(i).sizes()) + want = [TINY.max_prefill if d == "S" else d for d in tensor["shape"]] + self.assertEqual(sizes, want, f"{method} input {i}") + self.assertEqual(module.run_method("sample_rate", [])[0], TINY.sample_rate) + eager = build_programs(model) + inputs = example_inputs(TINY) + with torch.no_grad(): + for method in contract: + want = eager[method](*inputs[method]) + want = want if isinstance(want, tuple) else (want,) + got = module.run_method(method, list(inputs[method])) + for a, b in zip(got, want): + self.assertLess(max_diff(a, b), 1e-4, method) + + def test_xnnpack_program_runs_from_file(self): + # The runners load the program from a file, which frees delegate data + # after init; a load from a buffer never frees it. + model, program, _, _ = self._lower("xnnpack") + try: + import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch, + ) + except ImportError: + self.skipTest("ExecuTorch pybindings are not available") + inputs = example_inputs(TINY)["mimi_decode"] + with torch.no_grad(): + want = build_programs(model)["mimi_decode"](*inputs) + with tempfile.TemporaryDirectory() as d: + with open(f"{d}/phonon.pte", "wb") as f: + program.write_to_file(f) + got = _load_for_executorch(f"{d}/phonon.pte").run_method( + "mimi_decode", list(inputs) + ) + for a, b in zip(got, want): + self.assertLess(max_diff(a, b), 1e-4) + + def test_bf16_only_on_supported_backends(self): + with self.assertRaises(ValueError): + lower_phonon(random_model().to_half_split_rope(), "vulkan", dtype="bf16") + + def test_quantize_only_with_a_kernel(self): + self.assertEqual(quantize_scheme("xnnpack", "4bit"), "8da4w") + self.assertIsNone(quantize_scheme("vulkan", None)) + with self.assertRaises(ValueError): + quantize_scheme("vulkan", "4bit") + + def test_xnnpack_lowering_delegates(self): + _, _, delegation, _ = self._lower("xnnpack") + for method in ("voice_embed", "flow_lm", "mimi_decode"): + self.assertGreater(delegation[method]["delegates"], 0, method) + + +class TestHost(unittest.TestCase): + def test_prepare_text(self): + self.assertEqual(prepare_text(" hello world "), ("Hello world.", 3)) + self.assertEqual( + prepare_text("one two three four five?"), ("One two three four five?", 1) + ) + + def test_prepare_text_uses_ascii_rules(self): + # Same rules as the C++ runner: only ASCII letters and digits count. + self.assertEqual(prepare_text("élan"), ("élan.", 3)) + self.assertEqual(prepare_text("café"), ("Café", 3)) + + def test_noise_generator(self): + # The first splitmix64 output for seed 0. + self.assertEqual( + NoiseGenerator(0)._uniform(), (0xE220A8397B1DCDAF >> 11) * 2.0**-53 + ) + noise = NoiseGenerator(3).normal(20000, 1.0) + self.assertLess(abs(noise.mean().item()), 0.03) + self.assertLess(abs(noise.std().item() - 1.0), 0.03) + self.assertTrue( + torch.equal(NoiseGenerator(3).normal(32, 0.5), noise[:32] * 0.5) + ) + self.assertFalse(torch.equal(NoiseGenerator(4).normal(32, 1.0), noise[:32])) + + def test_split_sentences(self): + count = lambda s: len(s.split()) # noqa: E731 + text = "One two three. Four five six seven. Eight." + self.assertEqual( + split_sentences(text, count, 4), + ["One two three.", "Four five six seven.", "Eight."], + ) + self.assertEqual(split_sentences(text, count, 100), [text]) + self.assertEqual(split_sentences("a b c d e", count, 2), ["a b", "c d", "e"]) + + def test_runner_on_eager_programs(self): + model = random_model().to_half_split_rope() + programs = EagerPrograms(build_programs(model), program_metadata(model)) + runner = PhononRunner(programs, lambda text: [1, 2, 3][: len(text.split())]) + speaker = torch.randn(1, 6, TINY.dim) + pcm, info = runner.synthesize("hi there", speaker, temperature=0.0) + self.assertGreater(info["frames"], 0) + self.assertLessEqual(info["frames"], frame_budget(2, TINY.frame_rate)) + self.assertEqual(pcm.size, info["frames"] * TINY.frame_samples) + again, _ = runner.synthesize("hi there", speaker, temperature=0.0) + self.assertTrue((pcm == again).all()) + + def test_contract_is_json_serializable(self): + json.dumps(io_contract(PhononConfig())) + self.assertEqual(len(io_contract(PhononConfig())["mimi_decode"]["inputs"]), 16) + + +if __name__ == "__main__": + unittest.main() From 8ba612643ca0071c725c3db75940959205fb1624 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 15:57:59 -0800 Subject: [PATCH 11/15] Phonon: streaming output and per-backend weight quantization (#24929) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24929 Follow-up to D128116963. The C++ runner gets `--stream`, which writes 16-bit mono PCM to stdout as each 80 ms frame is decoded instead of writing the WAV at the end, so playback can start after the first frame (`phonon_runner ... --stream | aplay -q -f S16_LE -r 24000 -c 1`). The runner already decoded frame by frame through a callback; this only routes the samples to stdout and moves the summary line to stderr. `--qlinear` now accepts only the modes that each backend has a kernel for (XNNPACK 8da4w/8da8w, CUDA 4w/8w, Vulkan and WebGPU 4w/8da4w, MLX 4w/8w) and rejects the rest at export instead of failing at runtime. Quantization covers the flow LM's transformer and flow net; the input projection, the one-output EOS head and the Mimi codec stay in the export dtype. The EOS head is skipped because the Vulkan and WebGPU 4-bit linear does not yet handle an output size that is not a multiple of 4. CUDA 4w uses the packed int4 kernel, which needs `--dtype bf16` and packs on the GPU, so two `torch.arange` calls in the attention take the input's device. MLX 4w uses asymmetric int4: symmetric int4 keeps the words but delays the end-of-speech decision, adding up to two seconds of silence at the end of a sentence, while asymmetric int4 keeps the timing of FP32. The README documents both features with the sizes per backend. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: mergennachin Differential Revision: D131277773 --- examples/models/phonon/README.md | 29 +++++- examples/models/phonon/export_phonon.py | 108 +++++++++++++++----- examples/models/phonon/main.cpp | 62 ++++++++--- examples/models/phonon/model.py | 4 +- examples/models/phonon/tests/test_phonon.py | 30 +++++- 5 files changed, 183 insertions(+), 50 deletions(-) diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index 2c7f568ed3f..be74338d754 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -48,11 +48,34 @@ the same program (`python -m executorch.examples.models.phonon.run_phonon --model phonon_exports/$BACKEND/phonon.pte --text "Hello from ExecuTorch." --output hello_py.wav`, plus `--data-path` for CUDA). +`--stream` writes 16-bit mono PCM to stdout as each 80 ms frame is decoded, +instead of the WAV file, so playback starts after the first frame: + +```bash +cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --stream | + aplay -q -f S16_LE -r 24000 -c 1 # macOS: play -q -t raw -r 24000 -e signed -b 16 -c 1 - +``` + XNNPACK delegates the linear and convolution ops; the flow LM's attention runs as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. -`--quantize 4bit` (8-bit dynamic activations, 4-bit weights) or `--quantize 8bit` -quantizes the flow LM's linear layers on XNNPACK. A backend without a kernel for -the level fails before the download and names the options that work. + +`--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow +LM's transformer and flow net with the scheme of the backend; the codec stays in +the export dtype. Without `--quantize`, the export is full precision. + +| Backend | `--quantize 4bit` | `--quantize 8bit` | Size of phonon.pte (FP32: about 155 MiB) | +|---------|-------------------|-------------------|------------------------------------------| +| XNNPACK | 8-bit dynamic activations, 4-bit weights | 8-bit dynamic activations, 8-bit weights | 64, 75 MiB | +| CUDA | 4-bit weights, BF16 | 8-bit weights | 64, 62 MiB with BF16 | +| Vulkan, WebGPU | 4-bit weights | not available | 66 MiB | +| MLX | asymmetric 4-bit weights | 8-bit weights | 68, 76 MiB; 41, 51 MiB with BF16 | + +`--backend cuda --quantize 4bit` exports with `--dtype bf16`, which the int4 +kernel needs. A level without a kernel on the backend fails before the download +and names the options that work. 8-bit and the asymmetric 4-bit of CUDA and MLX +keep the timing of FP32. The symmetric 4-bit of XNNPACK, Vulkan and WebGPU +keeps the words but can hold the silence at the end of a sentence for up to two +seconds. Vulkan and MLX run each method in one partition. Vulkan needs `glslc`, for example from the [Vulkan SDK](https://vulkan.lunarg.com/sdk/home) (on Linux, diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index 0bf5f5181fa..7784fc65d33 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -54,8 +54,16 @@ # Backends that support --dtype bf16. BF16_BACKENDS = ("portable", "cuda", "mlx") # The linear quantization scheme of each --quantize level, per backend: 8da is -# dynamically quantized 8-bit activations with 4- or 8-bit weights. -QUANTIZE_SCHEMES = {"xnnpack": {"4bit": "8da4w", "8bit": "8da8w"}} +# dynamically quantized 8-bit activations, 4w and 8w quantize the weights only. +# CUDA 4w is the packed int4 kernel, which takes BF16 weights; MLX 4w is +# asymmetric (see quantize_asymmetric_int4). +QUANTIZE_SCHEMES = { + "xnnpack": {"4bit": "8da4w", "8bit": "8da8w"}, + "cuda": {"4bit": "4w", "8bit": "8w"}, + "vulkan": {"4bit": "4w"}, + "webgpu": {"4bit": "4w"}, + "mlx": {"4bit": "4w", "8bit": "8w"}, +} def attention_for(backend: str) -> str: @@ -83,12 +91,19 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): }.get(backend, []) -def quantize_scheme(backend: str, quantize: str | None) -> str | None: - """The linear quantization scheme for `--quantize` on `backend`. Raises - ValueError with what to run instead when the backend has no kernel for the - level.""" +def quantize_scheme( + backend: str, quantize: str | None, dtype: str | None +) -> tuple[str | None, str]: + """The linear quantization scheme and the dtype for `--quantize` and + `--dtype` on `backend`. Raises ValueError with what to run instead when the + backend has no kernel for them.""" + if dtype == "bf16" and backend not in BF16_BACKENDS: + raise ValueError( + f"--dtype bf16 is available for --backend {', '.join(BF16_BACKENDS)}. " + f"Export --backend {backend} without --dtype (FP32)." + ) if quantize is None: - return None + return None, dtype or "fp32" schemes = QUANTIZE_SCHEMES.get(backend, {}) if quantize not in schemes: options = [f"with --quantize {q}" for q in schemes] @@ -97,7 +112,35 @@ def quantize_scheme(backend: str, quantize: str | None) -> str | None: f"--backend {backend} has no {quantize} linear kernel. Export " f"{' or '.join(options)} instead." ) - return schemes[quantize] + scheme = schemes[quantize] + if backend == "cuda" and scheme == "4w": + if dtype == "fp32": + raise ValueError( + "--backend cuda --quantize 4bit needs BF16 weights (the int4 " + "kernel takes BF16). Drop --dtype fp32 and it uses BF16, or " + "export without --quantize for FP32." + ) + if dtype is None: + print("--backend cuda --quantize 4bit: using --dtype bf16") + dtype = "bf16" + return scheme, dtype or "fp32" + + +def quantize_asymmetric_int4(module: torch.nn.Module, group_size: int) -> None: + from torchao.quantization.granularity import PerGroup + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + from torchao.quantization.quant_primitives import MappingType + + quantize_( + module, + IntxWeightOnlyConfig( + weight_dtype=torch.int4, + granularity=PerGroup(group_size), + mapping_type=MappingType.ASYMMETRIC, + ), + filter_fn=lambda m, _: isinstance(m, torch.nn.Linear) + and m.in_features % group_size == 0, + ) def partitioners_for(backend: str, methods, quantized: bool): @@ -192,7 +235,7 @@ def lower_phonon( quantize: str | None = None, quantize_group_size: int | None = None, webgpu_compat_ops: bool = False, - dtype: str = "fp32", + dtype: str | None = None, ): """Export and lower the four programs of `model` for `backend`. Returns (ExecutorchProgramManager, delegation summary, metadata).""" @@ -201,12 +244,8 @@ def lower_phonon( ) cfg = model.cfg - qlinear = quantize_scheme(backend, quantize) - if dtype == "bf16" and (backend not in BF16_BACKENDS or qlinear): - raise ValueError( - f"--dtype bf16 is available for --backend {', '.join(BF16_BACKENDS)} " - "without --quantize. Export without --dtype (FP32)." - ) + qlinear, dtype = quantize_scheme(backend, quantize, dtype) + packing = "tile_packed_to_4d" if backend == "cuda" and qlinear == "4w" else None # Before the module swaps, which change the model's modules in place. metadata = program_metadata(model) programs = build_programs(model, attention_for(backend)) @@ -219,15 +258,27 @@ def lower_phonon( if qlinear: from executorch.extension.llm.export.quantize import quantize_model_ - # Weight quantization applies to the flow LM; the codec stays in fp32. - quantize_model_( - programs["flow_lm"], - qlinear_config=qlinear, - qlinear_group_size=quantize_group_size, - skip_incompatible_shapes=True, - ) + # Weight quantization applies to the transformer and flow net of the + # flow LM. The codec, the input projection and the one-output EOS head + # are not quantized. The int4 packing for CUDA runs on the GPU. + flow_lm = programs["flow_lm"].to("cuda" if packing else "cpu") + for module in (flow_lm.layers, flow_lm.flow_net): + if backend == "mlx" and qlinear == "4w": + # Symmetric int4 delays the end-of-speech decision (long + # pauses); asymmetric int4 keeps the timing of FP32. + quantize_asymmetric_int4(module, quantize_group_size or 32) + continue + quantize_model_( + module, + qlinear_config=qlinear, + qlinear_group_size=quantize_group_size, + qlinear_packing_format=packing, + skip_incompatible_shapes=True, + ) inputs, shapes = example_inputs(cfg), dynamic_shapes(cfg) + if packing: + inputs["flow_lm"] = tuple(t.to("cuda") for t in inputs["flow_lm"]) partitioners, passes, decompositions = partitioners_for( backend, list(programs), quantized=bool(qlinear) ) @@ -266,11 +317,12 @@ def export_phonon( quantize: str | None = None, quantize_group_size: int | None = None, webgpu_compat_ops: bool = False, - dtype: str = "fp32", + dtype: str | None = None, ) -> Path: """Export a checkpoint directory to output_dir/phonon.pte (plus aoti_cuda_blob.ptd for CUDA), phonon.json, the tokenizer and the default voice.""" + qlinear, dtype = quantize_scheme(backend, quantize, dtype) model = Phonon.from_checkpoint(checkpoint_dir / "model.safetensors") cfg = model.cfg program, delegation, metadata = lower_phonon( @@ -294,7 +346,7 @@ def export_phonon( "backend": backend, "dtype": dtype, "quantize": quantize, - "qlinear": quantize_scheme(backend, quantize), + "qlinear": qlinear, "webgpu_compat_ops": webgpu_compat_ops, "files": { "program": pte.name, @@ -336,8 +388,8 @@ def main(): parser.add_argument( "--dtype", choices=tuple(DTYPES), - default="fp32", - help="Weight and matmul dtype; inputs, outputs and state stay float32", + help="Weight and matmul dtype (default: fp32, or bf16 for --backend cuda " + "--quantize 4bit); inputs, outputs and state stay float32", ) parser.add_argument( "--webgpu-compat-ops", @@ -348,7 +400,7 @@ def main(): args = parser.parse_args() # Before the download, so that an unsupported combination fails first. try: - quantize_scheme(args.backend, args.quantize) + _, dtype = quantize_scheme(args.backend, args.quantize, args.dtype) except ValueError as e: parser.error(str(e)) export_phonon( @@ -358,7 +410,7 @@ def main(): quantize=args.quantize, quantize_group_size=args.quantize_group_size, webgpu_compat_ops=args.webgpu_compat_ops, - dtype=args.dtype, + dtype=dtype, ) diff --git a/examples/models/phonon/main.cpp b/examples/models/phonon/main.cpp index 59a53af3026..7369fea8007 100644 --- a/examples/models/phonon/main.cpp +++ b/examples/models/phonon/main.cpp @@ -10,6 +10,11 @@ // // phonon_runner --model_path phonon_exports/xnnpack/phonon.pte \ // --text "Hello from ExecuTorch." --output hello.wav +// +// --stream plays the audio while it is generated: +// +// phonon_runner --model_path ... --text "..." --stream | +// aplay -q -f S16_LE -r 24000 #include "runner.h" @@ -39,6 +44,11 @@ DEFINE_string( "program)"); DEFINE_string(text, "", "Text to synthesize"); DEFINE_string(output, "phonon.wav", "Output WAV file (16-bit PCM, mono)"); +DEFINE_bool( + stream, + false, + "Write raw 16-bit little-endian mono PCM to stdout as each frame is " + "decoded, instead of the WAV file"); DEFINE_double( temperature, -1, @@ -67,7 +77,14 @@ void write_u16(std::ofstream& out, uint16_t v) { out.write(reinterpret_cast(b), 2); } -// 16-bit mono PCM, rounded like run_phonon.py (round half to even). +// Rounded like run_phonon.py (round half to even). +int16_t to_pcm16(float sample) { + const double value = std::isfinite(sample) + ? std::nearbyint(static_cast(sample) * 32767.0) + : 0.0; + return static_cast(std::clamp(value, -32768.0, 32767.0)); +} + void write_wav( const std::string& path, const std::vector& pcm, @@ -87,11 +104,7 @@ void write_wav( out.write("data", 4); write_u32(out, data_bytes); for (float sample : pcm) { - double value = std::isfinite(sample) - ? std::nearbyint(static_cast(sample) * 32767.0) - : 0.0; - value = std::clamp(value, -32768.0, 32767.0); - write_u16(out, static_cast(static_cast(value))); + write_u16(out, static_cast(to_pcm16(sample))); } if (!out) { throw std::runtime_error("Could not write " + path); @@ -130,24 +143,47 @@ int main(int argc, char** argv) { : FLAGS_voice_path); const double load_s = seconds_since(start); + double first_audio_s = 0; + std::vector bytes; + const auto on_audio = [&](const float* pcm, size_t size) { + if (first_audio_s == 0) { + first_audio_s = seconds_since(start); + } + if (!FLAGS_stream) { + return; + } + bytes.resize(size * 2); + for (size_t i = 0; i < size; ++i) { + const auto v = static_cast(to_pcm16(pcm[i])); + bytes[2 * i] = static_cast(v); + bytes[2 * i + 1] = static_cast(v >> 8); + } + std::fwrite(bytes.data(), 1, bytes.size(), stdout); + std::fflush(stdout); + }; start = std::chrono::steady_clock::now(); - const auto result = - runner.synthesize(FLAGS_text, voice, FLAGS_temperature, FLAGS_seed); + const auto result = runner.synthesize( + FLAGS_text, voice, FLAGS_temperature, FLAGS_seed, on_audio); const double synth_s = seconds_since(start); if (result.pcm.empty()) { throw std::runtime_error("No audio generated"); } - write_wav(FLAGS_output, result.pcm, runner.sample_rate()); + if (!FLAGS_stream) { + write_wav(FLAGS_output, result.pcm, runner.sample_rate()); + } const double audio_s = static_cast(result.pcm.size()) / runner.sample_rate(); - std::printf( - "Wrote %s: %.2f s of audio, %lld frames, load %.2f s, synthesis " - "%.2f s (RTF %.2f)\n", - FLAGS_output.c_str(), + std::fprintf( + FLAGS_stream ? stderr : stdout, + "%s %s: %.2f s of audio, %lld frames, load %.2f s, first audio " + "%.3f s, synthesis %.2f s (RTF %.2f)\n", + FLAGS_stream ? "Streamed" : "Wrote", + FLAGS_stream ? "to stdout" : FLAGS_output.c_str(), audio_s, static_cast(result.frames), load_s, + first_audio_s, synth_s, synth_s / audio_s); } catch (const std::exception& error) { diff --git a/examples/models/phonon/model.py b/examples/models/phonon/model.py index 9701bf4d892..dbc58b668f4 100644 --- a/examples/models/phonon/model.py +++ b/examples/models/phonon/model.py @@ -736,12 +736,12 @@ def _attend_custom(self, i, q, k, v, input_pos): def _attend_sdpa(self, i, q, k, v, input_pos): cfg = self.cfg seq = q.shape[1] - pos = input_pos + torch.arange(seq, dtype=torch.long) + pos = input_pos + torch.arange(seq, dtype=torch.long, device=input_pos.device) q, k = hf_apply_rotary_emb(q, k, self.rope_cos[pos], self.rope_sin[pos]) k_cache, v_cache = getattr(self, f"k_cache_{i}"), getattr(self, f"v_cache_{i}") k_cache.index_copy_(2, pos, k.transpose(1, 2)) v_cache.index_copy_(2, pos, v.transpose(1, 2)) - keys = torch.arange(cfg.max_context, dtype=torch.long) + keys = torch.arange(cfg.max_context, dtype=torch.long, device=pos.device) mask = torch.where(keys[None, :] <= pos[:, None], 0.0, float("-inf")) o = F.scaled_dot_product_attention( q.transpose(1, 2), k_cache, v_cache, attn_mask=mask diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index 5553d4ee4ed..ce952bfec70 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -20,6 +20,7 @@ from executorch.examples.models.phonon.export_phonon import ( lower_phonon, + quantize_asymmetric_int4, quantize_scheme, source_transforms_for, ) @@ -341,11 +342,32 @@ def test_bf16_only_on_supported_backends(self): with self.assertRaises(ValueError): lower_phonon(random_model().to_half_split_rope(), "vulkan", dtype="bf16") + def test_quantize_picks_the_scheme_of_the_backend(self): + self.assertEqual(quantize_scheme("xnnpack", "4bit", None), ("8da4w", "fp32")) + self.assertEqual(quantize_scheme("vulkan", "4bit", None), ("4w", "fp32")) + self.assertEqual(quantize_scheme("cuda", "4bit", None), ("4w", "bf16")) + self.assertEqual(quantize_scheme("mlx", "8bit", "bf16"), ("8w", "bf16")) + self.assertEqual(quantize_scheme("vulkan", None, None), (None, "fp32")) + def test_quantize_only_with_a_kernel(self): - self.assertEqual(quantize_scheme("xnnpack", "4bit"), "8da4w") - self.assertIsNone(quantize_scheme("vulkan", None)) - with self.assertRaises(ValueError): - quantize_scheme("vulkan", "4bit") + for backend, quantize, dtype in ( + ("vulkan", "8bit", None), + ("portable", "4bit", None), + ("cuda", "4bit", "fp32"), + ("xnnpack", None, "bf16"), + ): + with self.assertRaises(ValueError, msg=f"{backend} {quantize}"): + quantize_scheme(backend, quantize, dtype) + + def test_asymmetric_int4_skips_incompatible_linears(self): + torch.manual_seed(0) + module = torch.nn.Sequential(torch.nn.Linear(64, 16), torch.nn.Linear(16, 8)) + x = torch.randn(4, 64) + want = module(x) + quantize_asymmetric_int4(module, 32) + self.assertIsNot(type(module[0].weight), torch.nn.Parameter) + self.assertIs(type(module[1].weight), torch.nn.Parameter) + self.assertLess(max_diff(module(x), want), 0.1) def test_xnnpack_lowering_delegates(self): _, _, delegation, _ = self._lower("xnnpack") From 423496d8f742830f82244da2fba4c6a590b80441 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 15:58:04 -0800 Subject: [PATCH 12/15] Fix the linear_q4gsw reference implementation for output channels not a multiple of 4 (#24930) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24930 make_linear_q4gsw_op and make_linear_dq8ca_q4gsw_op pad the weight scales and the bias of a 4-bit linear to a multiple of 4 output channels, but they do not pad the packed weights. The CompositeExplicitAutograd implementation of et_vk.linear_q4gsw (also used by et_vk.linear_dq8ca_q4gsw) dequantized the unpadded weights with the padded scales, so any lowering that re-traced the graph failed for such a linear, for example the 512 -> 1 end-of-speech head of a TTS model: RuntimeError: shape '[1, 16, 1]' is invalid for input of size 64 The implementation now drops the padded scale and bias rows. This affects the export-time reference implementation only; the dynamically quantized 8-bit activation, 4-bit weight kernel already handles any number of output channels on device. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: SS-JIA Differential Revision: D131277774 --- backends/vulkan/custom_ops_lib.py | 8 ++++-- backends/vulkan/test/test_vulkan_passes.py | 30 ++++++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index b674c6306ef..bc54ae6f66c 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -308,8 +308,12 @@ def linear_q4gsw( # Unpack the packed weights weights = unpack_4bit_weight_tensor(weights, x) - # Un-transpose the weight scales - weight_scales = weight_scales.transpose(0, 1) + # Un-transpose the weight scales. The scales and the bias are padded to a + # multiple of 4 output channels; the packed weights are not. + out_channels = weights.shape[0] + weight_scales = weight_scales.transpose(0, 1)[:out_channels] + if bias is not None: + bias = bias[:out_channels] weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32) weights = torch.ops.torchao.dequantize_affine( diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index 49dfb13af4c..cfaaeb653b1 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -1335,6 +1335,36 @@ def forward(self, x, y): ) self.assertTrue(torch.allclose(emb_x + emb_y, eager_ref, atol=1e-3, rtol=1e-3)) + def test_linear_q4gsw_reference_with_padded_scales_and_bias(self): + """make_linear_q4gsw_op pads the weight scales and the bias to a multiple + of 4 output channels, but not the packed weights. The reference + implementation must still return the unpadded output. + """ + import executorch.backends.vulkan.custom_ops_lib # noqa: registers et_vk ops + import torch.nn.functional as F + from executorch.backends.vulkan.patterns.weight_packing_utils import ( + pack_4bit_weight_tensor, + ) + + torch.manual_seed(0) + out_channels, in_channels, group_size = 6, 64, 32 + weight = torch.randint(-8, 8, (out_channels, in_channels), dtype=torch.int8) + scales = torch.rand(out_channels, in_channels // group_size) + 0.5 + bias = torch.randn(out_channels) + x = torch.randn(2, in_channels) + dq_weight = weight.float() * scales.repeat_interleave(group_size, dim=1) + expected = F.linear(x, dq_weight, bias) + + out = torch.ops.et_vk.linear_q4gsw.default( + x, + pack_4bit_weight_tensor(weight), + F.pad(scales.t(), (0, 2)), + group_size, + F.pad(bias, (0, 2)), + ) + self.assertEqual(out.shape, expected.shape) + self.assertTrue(torch.allclose(out, expected, atol=1e-4, rtol=1e-4)) + def test_register_param_mutation(self): """utils.register_param_mutation is a storage-keyed idempotency guard: the first call for a param returns True (proceed and record the tag), a From 518777b1de86c06fee196dbaaa3ba523ffa3e456 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 15:58:09 -0800 Subject: [PATCH 13/15] Do not delegate quantized linears that have no Vulkan kernel (#24971) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24971 The quantized linear pattern detector matched every linear with a dequantized constant weight, while replace_quantized_linear_patterns only has custom ops for 8-bit statically quantized activations and for 4-bit group-quantized weights (weight-only or with dynamically quantized activations). The partitioner delegates every node of a match, so an 8-bit weight-only linear (torchao IntxWeightOnlyConfig int8) was delegated with its torchao.dequantize_affine node left in place. The Vulkan runtime has no dequantize_affine kernel (the quantize/dequantize op implementations were removed in #15750), so the program failed to initialize on device. The registry entries for torchao.quantize_affine and torchao.dequantize_affine still claimed both ops as supported, which also delegated them outside of patterns. The detector now reports a match only when a custom op exists for it and the runtime can run it. The weight-only 4-bit kernel packs the weights in blocks of 4 output channels (Q4gswLinear.cpp checks N % 4 == 0 at prepack), so weight-only 4-bit linears with other output sizes are no longer matched; previously they exported and then failed at init. The match now also records the activation's quantize node in the dynamically quantized pattern, which the partitioner previously picked up only through the stale quantize_affine registry entry, and the two torchao registry entries are removed. Unmatched linears now stay outside of the delegate. With torchao quantization this surfaces at export as "Missing out variants: {'torchao::dequantize_affine'}" instead of a failure on device. WebGPU uses the same patterns, so the same applies there. Not changed: the quantized_decomposed quantize_per_channel / quantize_per_token / dequantize_per_channel / dequantize_per_token and choose_qparams (tensor, per_token_asymmetric) registry entries also have no runtime kernel since #15750. I left them because the VulkanQuantizer dynamic flow may still depend on them being claimed; they are worth a look by the backend owners. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: SS-JIA Differential Revision: D131277775 --- backends/vulkan/op_registry.py | 13 ---- backends/vulkan/patterns/quantized_linear.py | 36 +++++++++- backends/vulkan/test/test_vulkan_passes.py | 69 ++++++++++++++++++++ 3 files changed, 104 insertions(+), 14 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index b034e0ad0e5..3f9c550b779 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -568,19 +568,6 @@ def register_quantizedequantize_cpp_ops(): ) -@update_features( - [ - exir_ops.edge.torchao.quantize_affine.default, - exir_ops.edge.torchao.dequantize_affine.default, - ] -) -def register_torchao_quantize_dequantize(): - return OpFeatures( - inputs_storage=utils.CONTIGUOUS_BUFFER, - supports_resize=True, - ) - - # ============================================================================= # Q8taQuantizeDequantize.cpp # ============================================================================= diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index 0003c995810..353f1027aa8 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/backends/vulkan/patterns/quantized_linear.py @@ -157,6 +157,7 @@ def __init__(self, mm_node: torch.fx.Node) -> None: # noqa: C901 self.input_scales_node ): self.quantize_input_node = input_to_dq_node + self.all_nodes.append(self.quantize_input_node) self.pattern_input_node = self.quantize_input_node.args[0] # The implementation has a limitation that input channels must be a @@ -267,6 +268,39 @@ def is_weight_perchannel_quantized(self) -> bool: # scales should have same size as weight's output channels dim return scales_shape[0] == weight_shape[-2] + def is_weight_4bit_quantized(self) -> bool: + # Read the quant range from the dequantize node's arguments, since the + # detector has no access to the weight data. + node = self.dequantize_weight_node + arg_names = [arg.name for arg in node.target._schema.arguments] + args = {**dict(zip(arg_names, node.args)), **node.kwargs} + quant_min, quant_max = args.get("quant_min"), args.get("quant_max") + if quant_min is None or quant_max is None: + return False + return quant_min >= -8 and quant_max <= 7 + + def has_fused_op(self) -> bool: + """Whether this match becomes an op the runtime can run, through + replace_quantized_linear_patterns or FuseQuantizedOpsTransform. The + partitioner delegates every node of a match, so a match without one + would leave its dequantize nodes, which have no Vulkan implementation, + in the delegate.""" + if self.is_input_static_per_tensor_quantized(): + return self.is_weight_perchannel_quantized() + if self.is_weight_only_quantized() and utils.is_dequant_per_channel_node( + self.dequantize_weight_node + ): + # FuseQuantizedOpsTransform: _weight_int8pack_mm or linear_qcs4w. + return self.is_weight_perchannel_quantized() + if not ( + self.is_weight_pergroup_quantized() and self.is_weight_4bit_quantized() + ): + return False + if self.is_weight_only_quantized(): + # linear_q4gsw packs the weights in blocks of 4 output channels. + return self.weight_node.meta["val"].shape[-2] % 4 == 0 + return self.is_input_dynamic_perchannel_quantized() + def is_input_static_per_tensor_quantized(self) -> bool: if self.dequantize_input_node is None: return False @@ -313,7 +347,7 @@ def find_quantized_linear_patterns( return None matched_pattern = QuantizedLinearMatch(node) - if matched_pattern.match_found: + if matched_pattern.match_found and matched_pattern.has_fused_op(): return matched_pattern return None diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index cfaaeb653b1..11e8872a8d2 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -1503,6 +1503,75 @@ def _reshape_for_broadcast(self, freqs_cis: torch.Tensor, x: torch.Tensor): # We expect at least one custom op to be created self.assertGreater(custom_op_count, 0) + def test_quantized_linear_matches_only_runnable_linears(self): + """The quantized linear pattern matches only linears that have a Vulkan + custom op: 4-bit weights, with output channels in multiples of 4 for the + weight-only op. Other matches would be delegated without a kernel. + """ + from executorch.backends.vulkan.patterns.quantized_linear import ( + find_quantized_linear_patterns, + linear_anchor_nodes, + ) + from torchao.quantization.granularity import PerAxis, PerGroup + from torchao.quantization.quant_api import ( + Int8DynamicActivationIntxWeightConfig, + IntxWeightOnlyConfig, + quantize_, + ) + from torchao.utils import unwrap_tensor_subclass + + def matches(config, out_channels: int) -> List[bool]: + model = torch.nn.Sequential(torch.nn.Linear(64, out_channels)) + quantize_(model, config) + unwrap_tensor_subclass(model) + program = torch.export.export(model, (torch.randn(2, 64),), strict=True) + ep = to_edge( + program, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).exported_program() + return [ + find_quantized_linear_patterns(n) is not None + for n in ep.graph.nodes + if n.target in linear_anchor_nodes + ] + + w4 = IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)) + da8w4 = Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(32) + ) + w8 = IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0)) + self.assertEqual(matches(w4, 8), [True]) + self.assertEqual(matches(w4, 6), [False]) + self.assertEqual(matches(da8w4, 6), [True]) + self.assertEqual(matches(w8, 8), [False]) + + def test_partitioner_skips_quant_ops_without_runtime_kernel(self): + """The Vulkan runtime implements torchao.dequantize_affine only inside + fused patterns. An 8-bit weight-only linear has no fused pattern, so its + dequantize_affine must stay outside of the Vulkan delegate. + """ + from executorch.backends.vulkan.partitioner.vulkan_partitioner import ( + VulkanPartitioner, + ) + from executorch.exir import to_edge_transform_and_lower + from torchao.quantization.granularity import PerAxis + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + from torchao.utils import unwrap_tensor_subclass + + model = torch.nn.Sequential(torch.nn.Linear(64, 64)) + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0)), + ) + unwrap_tensor_subclass(model) + sample_inputs = (torch.randn(2, 64),) + edge_program = to_edge_transform_and_lower( + torch.export.export(model, sample_inputs, strict=True), + partitioner=[VulkanPartitioner()], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + gm = edge_program.exported_program().graph_module + self.assertEqual(op_node_count(gm, "dequantize_affine.default"), 1) + def test_fuse_q8ta_linear(self): """Test that sequential quantized linears fuse into q8ta_linear when output quantization is present.""" from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( From 3495700870fc4679f67ccf90dcd310f30bd3c5d4 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Thu, 12 Nov 2026 15:58:13 -0800 Subject: [PATCH 14/15] Phonon: run the Mimi decoder's transposed convolutions as convolutions on XNNPACK Summary: XNNPACK's deconvolution is the most expensive op in Phonon's Mimi decoder: about 45% of each frame on x86 and 38% on arm64. Every transposed convolution in Mimi has kernel = 2 * stride, so each one is exactly a kernel-2 convolution with stride * out_channels outputs, one per output phase, whose phases are then interleaved into the output steps. The XNNPACK export now applies that rewrite (conv_transpose1d_as_conv1d) to the three decoder transposed convolutions, and the existing depthwise_conv_transpose_as_mul to the upsampler. The output is the same up to float rounding: the exported 31M program writes audio within 1 LSB (16-bit PCM) of the program without the rewrite, with the same length and frames. Running the transposed convolutions as 2-D convolutions with the sequence on the width axis (conv_transpose1d_as_2d) does not help: Mimi stays at 8.2 ms per frame, so the gain comes from the convolution kernel rather than the tiling axis. This diff was written with the help of Claude Code (Claude Opus 5.5). Differential Revision: D131605676 --- examples/models/phonon/README.md | 4 +- examples/models/phonon/export_phonon.py | 1 + .../models/phonon/source_transformations.py | 57 ++++++++++++++++++- examples/models/phonon/tests/test_phonon.py | 33 +++++++++-- 4 files changed, 88 insertions(+), 7 deletions(-) diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index be74338d754..7f471832a30 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -57,7 +57,9 @@ cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --s ``` XNNPACK delegates the linear and convolution ops; the flow LM's attention runs -as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. The export also +rewrites the Mimi decoder's transposed convolutions as convolutions, which +XNNPACK runs faster. `--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow LM's transformer and flow net with the scheme of the backend; the codec stays in diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index 7784fc65d33..2bae25d8f47 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -82,6 +82,7 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") return [st.use_webgpu_compat_ops] return { + "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d], "mlx": [st.decompose_elu], "vulkan": [ st.decompose_elu, diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py index 69a2c228a54..4c6a1dcdb9b 100644 --- a/examples/models/phonon/source_transformations.py +++ b/examples/models/phonon/source_transformations.py @@ -4,7 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Function-preserving module swaps for backends that lack an op. +"""Function-preserving module swaps for backends that lack an op or run it +slowly. The model in model.py uses plain PyTorch ops. Each swap below computes the same function from more widely supported ops; export_phonon.py selects them @@ -13,10 +14,15 @@ decompose_elu MLX and Vulkan have no aten.elu kernel; without it the Mimi decoder splits at every ELU. conv_transpose1d_as_2d Vulkan runs transposed convolution on 4-D inputs only. + conv_transpose1d_as_conv1d + XNNPACK's deconvolution is slow; this runs the Mimi + decoder's transposed convolutions (kernel = + 2 * stride) as a convolution and an interleave. depthwise_conv_transpose_as_mul Vulkan computes the Mimi upsampler (depthwise, kernel = 2 * stride) incorrectly; this runs it as a - multiply and an overlap-add. + multiply and an overlap-add. XNNPACK also runs it + faster this way. use_webgpu_compat_ops WebGPU runtimes built before these runtime fixes (off by default; --webgpu-compat-ops): * native_layer_norm resize hook (the flow LM's @@ -93,6 +99,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y.squeeze(2) +class ConvTranspose1dAsConv1d(nn.Module): + """ConvTranspose1d with kernel_size == 2 * stride as a kernel-2 Conv1d + with stride * out_channels outputs (one per output phase), interleaved + into the output steps.""" + + def __init__(self, convtr: nn.ConvTranspose1d): + super().__init__() + (self.stride,) = convtr.stride + s = self.stride + w = convtr.weight.detach() # [cin, cout, 2s] + cin, cout, _ = w.shape + # Phase j of output step p reads input p - 1 through tap j + s and + # input p through tap j. + taps = torch.stack([w[..., s:], w[..., :s]], -1) # [cin, cout, s, 2] + self.weight = nn.Parameter( + taps.permute(1, 2, 0, 3).reshape(cout * s, cin, 2).contiguous(), + requires_grad=False, + ) + self.bias = convtr.bias + self.out_channels = cout + + def forward(self, x: torch.Tensor) -> torch.Tensor: + z = F.conv1d(x, self.weight, padding=1) # [b, cout * s, n + 1] + b, _, n = z.shape + s = self.stride + y = z.view(b, self.out_channels, s, n).transpose(2, 3) + y = y.reshape(b, self.out_channels, n * s) + return y if self.bias is None else y + self.bias.view(1, -1, 1) + + class ComputeIn(nn.Module): """Runs `module` with its weights in `dtype` on a float32 input and returns float32. Other attributes (bias, kernel_size, ...) are read from @@ -171,6 +207,23 @@ def conv_transpose1d_as_2d(programs: dict[str, nn.Module]) -> int: return count +def conv_transpose1d_as_conv1d(programs: dict[str, nn.Module]) -> int: + count = 0 + for program in programs.values(): + for module in program.modules(): + if not isinstance(module, StreamingConvTranspose1d): + continue + convtr = module.convtr + if ( + isinstance(convtr, nn.ConvTranspose1d) + and convtr.groups == 1 + and convtr.kernel_size[0] == 2 * convtr.stride[0] + ): + module.convtr = ConvTranspose1dAsConv1d(convtr) + count += 1 + return count + + def depthwise_conv_transpose_as_mul(programs: dict[str, nn.Module]) -> int: count = 0 for program in programs.values(): diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index ce952bfec70..63b79dee939 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -43,6 +43,7 @@ ) from executorch.examples.models.phonon.source_transformations import ( compute_weights_in, + ConvTranspose1dAsConv1d, use_webgpu_compat_ops, ) from executorch.examples.models.phonon.tests.reference import ( @@ -236,19 +237,43 @@ def test_webgpu_compat_ops_preserve_outputs(self): for a, b in zip(outputs, programs[name](*inputs[name])): self.assertLess(max_diff(a, b), 1e-4) - def test_vulkan_swaps_preserve_outputs(self): + def _check_swaps(self, backend: str, counts: list[int]): inputs = example_inputs(TINY) programs = build_programs(self.model) with torch.no_grad(): before = programs["mimi_decode"](*inputs["mimi_decode"]) - # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. - counts = [t(programs) for t in source_transforms_for("vulkan", False)] - self.assertEqual(counts, [7, 1, 2]) + self.assertEqual( + [t(programs) for t in source_transforms_for(backend, False)], counts + ) with torch.no_grad(): after = programs["mimi_decode"](*inputs["mimi_decode"]) for a, b in zip(before, after): self.assertLess(max_diff(a, b), 1e-4) + def test_vulkan_swaps_preserve_outputs(self): + # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("vulkan", [7, 1, 2]) + + def test_xnnpack_swaps_preserve_outputs(self): + # The depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("xnnpack", [1, 2]) + + def test_conv_transpose_as_conv_matches(self): + torch.manual_seed(0) + for stride, steps, bias in ( + (1, 1, True), + (4, 16, False), + (5, 1, True), + (6, 16, True), + ): + convtr = torch.nn.ConvTranspose1d( + 6, 3, 2 * stride, stride=stride, bias=bias + ) + x = torch.randn(1, 6, steps) + with torch.no_grad(): + got = ConvTranspose1dAsConv1d(convtr)(x) + self.assertLess(max_diff(got, convtr(x)), 1e-5, stride) + def test_compute_weights_in_keeps_float32_io(self): inputs = example_inputs(TINY) programs = build_programs(self.model) From aa1d71732f8d409e488fbe7f9236c2171d91e2b3 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 13 Nov 2026 11:50:16 -0500 Subject: [PATCH 15/15] [XNNPACK] Lower dense transposed convs as one polyphase conv + depth to space XNNPACK's deconvolution is much slower than its convolution for the upsamplers of audio decoders and vocoders, transposed convs with kernel = 2 * stride. DecomposeConvTransposePass lowers a dense (groups == 1) 1-D or 2-D transposed conv whose kernel is a multiple of its stride as one stride-1 conv that computes every output phase, an interleave of the phases (depth_to_space_2d in 2-D, a view, permute and view in 1-D) and a slice for the padding. The Phonon XNNPACK export drops conv_transpose1d_as_conv1d, which did the same rewrite for the Mimi decoder before export. This change was authored with Claude Code (Claude Opus 5.5). --- backends/xnnpack/_passes/__init__.py | 4 + .../_passes/decompose_conv_transpose.py | 219 ++++++++++++ .../passes/test_decompose_conv_transpose.py | 318 ++++++++++++++++++ examples/models/phonon/README.md | 4 +- examples/models/phonon/export_phonon.py | 2 +- .../models/phonon/source_transformations.py | 6 +- examples/models/phonon/tests/test_phonon.py | 5 +- 7 files changed, 549 insertions(+), 9 deletions(-) create mode 100644 backends/xnnpack/_passes/decompose_conv_transpose.py create mode 100644 backends/xnnpack/test/passes/test_decompose_conv_transpose.py diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index 5db1599f1ab..09911ec6e7d 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -26,6 +26,9 @@ ) from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate +from executorch.backends.xnnpack._passes.decompose_conv_transpose import ( + DecomposeConvTransposePass, +) from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( DecomposeDepthwiseConvTransposePass, ) @@ -93,6 +96,7 @@ def __init__( FuseBatchNormPass, DecomposeBatchNorm, DecomposeDepthwiseConvTransposePass, + DecomposeConvTransposePass, FuseActivationPass, DecomposeConcatenate, RemoveGetItemPass, diff --git a/backends/xnnpack/_passes/decompose_conv_transpose.py b/backends/xnnpack/_passes/decompose_conv_transpose.py new file mode 100644 index 00000000000..13a017a6e9b --- /dev/null +++ b/backends/xnnpack/_passes/decompose_conv_transpose.py @@ -0,0 +1,219 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import cast, List, Optional, Tuple + +import executorch.backends.xnnpack.custom_ops_lib # noqa: F401 registers et_xnnpack ops +import torch +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + delete_constant_placeholder, + get_param_tensor, + is_param_node, +) +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult + + +class DecomposeConvTransposePass(XNNPACKPass): + """ + Lower a dense (groups == 1) 1-D or 2-D transposed conv whose kernel is a + multiple of its stride s > 1 as one stride-1 conv that computes every + output phase, a depth to space that interleaves the phases, and a slice + for the padding. + + With k = m * s taps per axis, output phase j of a stride-s transposed conv + is an m-tap correlation of the input (before the padding p is cropped): + + y[s*q + j] = sum_{u < m} x[q + u - (m - 1)] * w[s*(m - 1 - u) + j] + + The conv has padding m - 1 and s * out_channels outputs per axis, ordered + phase-major so that the 2-D interleave is the block-major + depth_to_space_2d. XNNPACK runs it on its convolution kernels, which are + much faster than its deconvolution; k = 2 * s is the usual upsampler of + audio decoders. + + Depthwise transposed convs are left to DecomposeDepthwiseConvTransposePass. + Grouped ones and those with dilation, output_padding > padding, unequal + 2-D strides, a dequantized weight or dynamic shapes stay on XNNPACK's own + node. + """ + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + # Plan every conv before rewriting any: the nodes created by + # _decompose have no meta["val"] until the retrace below. + plans = [] + for node in graph.nodes: + if ( + node.op == "call_function" + and node.target == exir_ops.edge.aten.convolution.default + ): + plan = self._plan(node) + if plan is not None: + plans.append((node, plan)) + for node, plan in plans: + self._decompose(graph, node, *plan) + + modified = bool(plans) + if modified: + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) + + def _plan( + self, node: torch.fx.Node + ) -> Optional[Tuple[int, List[int], List[int], List[int]]]: + ( + x, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) = node.args + if not transposed or groups != 1: + return None + if not is_param_node(self.exported_program, weight): + return None + if bias is not None and not is_param_node(self.exported_program, bias): + return None + x_shape = list(x.meta["val"].shape) + out_shape = list(node.meta["val"].shape) + if not all(isinstance(d, int) for d in x_shape + out_shape): + return None + kernel = list(weight.meta["val"].shape[2:]) + dims = len(kernel) + if dims not in (1, 2): + return None + # ATen broadcasts a single spatial value over every spatial dim. + stride, padding, dilation, output_padding = ( + p * dims if len(p) == 1 else p + for p in cast(List[List[int]], [stride, padding, dilation, output_padding]) + ) + s = stride[0] + if ( + s < 2 + or any(v != s for v in stride) + or any(k % s for k in kernel) + or any(d != 1 for d in dilation) + or any(o > p for o, p in zip(output_padding, padding)) + ): + return None + return s, padding, x_shape, out_shape + + def _decompose( + self, + graph: torch.fx.Graph, + node: torch.fx.Node, + s: int, + padding: List[int], + x_shape: List[int], + out_shape: List[int], + ) -> None: + x, weight_node, bias_node = node.args[0], node.args[1], node.args[2] + weight = get_param_tensor(self.exported_program, weight_node) + assert weight is not None + c_in, c_out, *kernel = weight.shape + dims = len(kernel) + taps = [k // s for k in kernel] + phases = s**dims + + # w[ci, co, s*t + j] per axis -> phase_weight[(j, co), ci, m - 1 - t] + tap_axes = [2 + 2 * i for i in range(dims)] + phase_axes = [3 + 2 * i for i in range(dims)] + phase_weight = ( + weight.reshape(c_in, c_out, *[d for m in taps for d in (m, s)]) + .flip(tap_axes) + .permute(*phase_axes, 1, 0, *tap_axes) + .reshape(phases * c_out, c_in, *taps) + .contiguous() + ) + phase_weight_node = self._constant_like( + graph, weight_node, f"_s{s}_phases", phase_weight + ) + phase_bias_node = None + if bias_node is not None: + bias = get_param_tensor(self.exported_program, bias_node) + assert bias is not None + phase_bias_node = self._constant_like( + graph, bias_node, f"_x{phases}", bias.repeat(phases) + ) + + def call(target, *args) -> torch.fx.Node: + return graph.create_node("call_function", target, args) + + with graph.inserting_before(node): + y = call( + exir_ops.edge.aten.convolution.default, + x, + phase_weight_node, + phase_bias_node, + [1] * dims, + [m - 1 for m in taps], + [1] * dims, + False, + [0] * dims, + 1, + ) + if dims == 2: + y = call(exir_ops.edge.et_xnnpack.depth_to_space_2d.default, y, s) + else: + batch, steps = x_shape[0], x_shape[2] + taps[0] - 1 + y = call( + exir_ops.edge.aten.view_copy.default, y, [batch, s, c_out, steps] + ) + y = call(exir_ops.edge.aten.permute_copy.default, y, [0, 2, 3, 1]) + y = call( + exir_ops.edge.aten.view_copy.default, y, [batch, c_out, steps * s] + ) + for axis, p in enumerate(padding): + if p: + y = call( + exir_ops.edge.aten.slice_copy.Tensor, + y, + 2 + axis, + p, + p + out_shape[2 + axis], + 1, + ) + node.replace_all_uses_with(y) + graph.erase_node(node) + for constant in (weight_node, bias_node): + if constant is not None and len(constant.users) == 0: + delete_constant_placeholder(self.exported_program, constant) + + def _constant_like( + self, + graph: torch.fx.Graph, + like: torch.fx.Node, + suffix: str, + data: torch.Tensor, + ) -> torch.fx.Node: + spec = next( + spec + for spec in self.exported_program.graph_signature.input_specs + if spec.arg.name == like.name + ) + with graph.inserting_before(like): + node = create_constant_placeholder( + exp_program=self.exported_program, + graph=graph, + kind=spec.kind, + persistent_buffer=spec.persistent, + name=like.name + suffix, + data=data, + ) + # create_constant_placeholder only sets meta["val"]; keep the + # external-constant tag so the new constant follows the original into + # its .ptd file. + if "custom" in like.meta: + node.meta["custom"] = dict(like.meta["custom"]) + return node diff --git a/backends/xnnpack/test/passes/test_decompose_conv_transpose.py b/backends/xnnpack/test/passes/test_decompose_conv_transpose.py new file mode 100644 index 00000000000..adb597a6198 --- /dev/null +++ b/backends/xnnpack/test/passes/test_decompose_conv_transpose.py @@ -0,0 +1,318 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from unittest import mock + +import executorch.backends.test.harness.stages as BaseStages +import executorch.backends.xnnpack.test.passes.test_decompose_depthwise_conv_transpose as dw_tests +import torch +from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack._passes.decompose_conv_transpose import ( + DecomposeConvTransposePass, +) +from executorch.backends.xnnpack._passes.decompose_depthwise_conv_transpose import ( + DecomposeDepthwiseConvTransposePass, +) +from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass +from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir._serialize import _deserialize_pte_binary +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.external_constants_pass import ( + delegate_external_constants_pass_unlifted, +) +from executorch.extension.flat_tensor.serialize.serialize import ( + _deserialize_to_flat_tensor, +) + + +class TestDecomposeConvTransposePass(unittest.TestCase): + PassStage = RunPasses([DecomposeConvTransposePass]) + + conv = "executorch_exir_dialects_edge__ops_aten_convolution_default" + view = "executorch_exir_dialects_edge__ops_aten_view_copy_default" + permute = "executorch_exir_dialects_edge__ops_aten_permute_copy_default" + slice = "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor" + d2s = "executorch_exir_dialects_edge__ops_et_xnnpack_depth_to_space_2d_default" + + def setUp(self): + torch._dynamo.reset() + torch.manual_seed(0) + + @staticmethod + def _deconv(dims, c_in, c_out, kernel, stride, **kwargs): + cls = torch.nn.ConvTranspose1d if dims == 1 else torch.nn.ConvTranspose2d + return cls(c_in, c_out, kernel, stride=stride, **kwargs) + + @staticmethod + def _transposed_convs(tester) -> int: + graph = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ).graph + return sum( + node.target == exir_ops.edge.aten.convolution.default and node.args[6] + for node in graph.nodes + ) + + def _run_pass(self, module, inputs, dynamic_shapes=None): + return ( + Tester(module.eval(), inputs, dynamic_shapes=dynamic_shapes) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + def _check_rewritten(self, module, inputs, counts, tol=1e-4): + tester = self._run_pass(module, inputs).check_count(counts) + self.assertEqual(self._transposed_convs(tester), 0) + tester.run_method_and_compare_outputs(atol=tol, rtol=tol) + return tester + + def _check_left_alone(self, module, inputs, dynamic_shapes=None): + tester = self._run_pass(module, inputs, dynamic_shapes) + self.assertEqual(self._transposed_convs(tester), 1) + + def _check_lowered(self, module, inputs, tol=1e-4, num_deconvs=1): + # XNNPACK delegates these transposed convs on their own, so one + # delegate and matching outputs do not prove the pass ran during + # lowering: count its calls and check the serialized graph. + tester = Tester(module.eval(), inputs).export() + with mock.patch.object( + DecomposeConvTransposePass, + "_decompose", + autospec=True, + side_effect=DecomposeConvTransposePass._decompose, + ) as decompose: + tester.to_edge_transform_and_lower() + self.assertEqual(decompose.call_count, num_deconvs) + tester.check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + node_types = ( + dw_tests.TestDecomposeDepthwiseConvTransposePass._xnnpack_node_types(tester) + ) + self.assertEqual(node_types["XNNConvTranspose2d"], 0, node_types) + self.assertGreaterEqual(node_types["XNNConv2d"], num_deconvs, node_types) + ( + tester.to_executorch() + .serialize() + .run_method_and_compare_outputs(num_runs=3, atol=tol, rtol=tol) + ) + + def test_1d_kernel_twice_stride(self): + for stride in (2, 4, 5, 6): + for bias in (True, False): + with self.subTest(stride=stride, bias=bias): + module = self._deconv(1, 6, 3, 2 * stride, stride, bias=bias) + inputs = (torch.randn(1, 6, 7),) + self._check_rewritten( + module, + inputs, + {self.conv: 1, self.view: 2, self.permute: 1, self.slice: 0}, + ) + self._check_lowered(module, inputs) + + def test_1d_single_step(self): + # Streaming decoders run one input step per call. + module = self._deconv(1, 8, 4, 8, 4) + inputs = (torch.randn(1, 8, 1),) + self._check_rewritten(module, inputs, {self.conv: 1, self.view: 2}) + self._check_lowered(module, inputs) + + def test_1d_padding_and_output_padding(self): + for kernel, stride, padding, output_padding in ( + (8, 4, 2, 1), + (16, 8, 4, 0), + (4, 2, 1, 1), + ): + with self.subTest(kernel=kernel, stride=stride, padding=padding): + module = self._deconv( + 1, + 5, + 4, + kernel, + stride, + padding=padding, + output_padding=output_padding, + ) + inputs = (torch.randn(2, 5, 6),) + self._check_rewritten(module, inputs, {self.conv: 1, self.slice: 1}) + self._check_lowered(module, inputs) + + def test_1d_kernel_multiple_of_stride(self): + # m = 1 needs no padding, m = 3 pads the conv by 2. + for kernel, stride in ((3, 3), (6, 2), (9, 3)): + with self.subTest(kernel=kernel, stride=stride): + module = self._deconv(1, 4, 6, kernel, stride) + inputs = (torch.randn(1, 4, 5),) + self._check_rewritten(module, inputs, {self.conv: 1, self.view: 2}) + self._check_lowered(module, inputs) + + def test_2d(self): + for kernel, stride, padding, output_padding, slices in ( + (4, 2, 1, 0, 2), + ((4, 2), 2, 0, 0, 0), + (6, 3, (1, 2), (1, 0), 2), + ): + with self.subTest(kernel=kernel, stride=stride, padding=padding): + module = self._deconv( + 2, + 4, + 6, + kernel, + stride, + padding=padding, + output_padding=output_padding, + ) + inputs = (torch.randn(1, 4, 5, 7),) + self._check_rewritten( + module, inputs, {self.conv: 1, self.d2s: 1, self.slice: slices} + ) + self._check_lowered(module, inputs) + + def test_fp16(self): + module = self._deconv(1, 8, 4, 8, 4).half() + inputs = (torch.randn(1, 8, 6).half(),) + self._check_rewritten(module, inputs, {self.conv: 1, self.view: 2}, tol=1e-2) + self._check_lowered(module, inputs, tol=1e-2) + + def test_two_deconvs_with_activation(self): + # The second deconv's input is the first one's interleave. + module = torch.nn.Sequential( + self._deconv(1, 16, 8, 12, 6), + torch.nn.ELU(), + self._deconv(1, 8, 4, 10, 5), + ) + inputs = (torch.randn(1, 16, 3),) + self._check_rewritten(module, inputs, {self.conv: 2, self.view: 4}) + self._check_lowered(module, inputs, num_deconvs=2) + + class SharedDeconv(torch.nn.Module): + def __init__(self): + super().__init__() + self.deconv = torch.nn.ConvTranspose1d(6, 4, 8, stride=4) + + def forward(self, x, z): + return self.deconv(x) + self.deconv(z) + + def test_shared_weight(self): + # Both deconvs reuse one phase weight and one phase bias. + module = self.SharedDeconv() + inputs = (torch.randn(1, 6, 5), torch.randn(1, 6, 5)) + tester = self._check_rewritten(module, inputs, {self.conv: 2, self.view: 4}) + placeholders = [ + node.name + for node in tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module.graph.nodes + if node.op == "placeholder" + ] + self.assertEqual(sum(n.endswith("_s4_phases") for n in placeholders), 1) + self.assertEqual(sum(n.endswith("_x4") for n in placeholders), 1) + + class DeconvBiasReused(torch.nn.Module): + """The overlap-add of a streaming decoder reads the bias again.""" + + def __init__(self): + super().__init__() + self.deconv = torch.nn.ConvTranspose1d(6, 4, 8, stride=4) + + def forward(self, x): + return self.deconv(x) - self.deconv.bias.view(1, -1, 1) + + def test_reused_bias(self): + module = self.DeconvBiasReused() + inputs = (torch.randn(1, 6, 5),) + # The interleave and the bias view. + self._check_rewritten(module, inputs, {self.conv: 1, self.view: 3}) + self._check_lowered(module, inputs) + + def test_pass_order(self): + module = self._deconv(1, 6, 3, 4, 2) + inputs = (torch.randn(1, 6, 7),) + passes = XNNPACKPassManager( + Tester(module.eval(), inputs).export().get_artifact() + ).passes + self.assertLess( + passes.index(DecomposeDepthwiseConvTransposePass), + passes.index(DecomposeConvTransposePass), + ) + self.assertLess( + passes.index(DecomposeConvTransposePass), + passes.index(FuseActivationPass), + ) + + class ExportWithExternalConstants(BaseStages.Export): + def run(self, artifact, inputs): + tagged = torch.export.export(artifact, inputs, strict=True).module() + delegate_external_constants_pass_unlifted( + module=tagged, gen_tag_fn=lambda _: "model" + ) + self.exported_program = torch.export.export(tagged, inputs, strict=True) + + def test_external_constants(self): + # The phase weight and bias inherit the external-constant tag of the + # constants they replace. + module = self._deconv(1, 6, 3, 8, 4).eval() + inputs = (torch.randn(1, 6, 7),) + tester = Tester(module, inputs) + tester.export(self.ExportWithExternalConstants()) + tester.to_edge_transform_and_lower().check_count( + {"torch.ops.higher_order.executorch_call_delegate": 1} + ).to_executorch() + program = tester.get_artifact() + self.assertIsNone(_deserialize_pte_binary(program.buffer).named_data) + ptd = _deserialize_to_flat_tensor(bytes(program._tensor_data["model"])) + self.assertEqual(len(ptd.named_data), 2) + + from executorch.extension.pybindings import portable_lib as runtime + + loaded = runtime._load_for_executorch_from_buffer( + program.buffer, bytes(program._tensor_data["model"]) + ) + self.assertTrue( + torch.allclose(loaded.forward(inputs)[0], module(*inputs), atol=1e-4) + ) + + def test_left_alone(self): + for name, module, shape in ( + ("depthwise", self._deconv(1, 4, 4, 4, 2, groups=4), (1, 4, 5)), + ("grouped", self._deconv(1, 4, 4, 4, 2, groups=2), (1, 4, 5)), + ("dilated", self._deconv(1, 4, 4, 4, 2, dilation=2), (1, 4, 5)), + ("kernel not a multiple", self._deconv(1, 4, 4, 3, 2), (1, 4, 5)), + ("stride 1", self._deconv(1, 4, 4, 2, 1), (1, 4, 5)), + ( + "output_padding > padding", + self._deconv(1, 4, 4, 4, 2, output_padding=1), + (1, 4, 5), + ), + ("unequal strides", self._deconv(2, 4, 4, (4, 2), (2, 1)), (1, 4, 5, 5)), + ): + with self.subTest(name): + self._check_left_alone(module, (torch.randn(*shape),)) + + def test_dynamic_shape_is_left_alone(self): + module = self._deconv(1, 4, 4, 4, 2) + self._check_left_alone( + module, + (torch.randn(1, 4, 6),), + dynamic_shapes=({2: torch.export.Dim("steps", min=2, max=16)},), + ) + + def test_dequantized_weight_is_left_alone(self): + class QuantizedDeconv(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "weight", torch.randint(-128, 127, (4, 4, 4), dtype=torch.int8) + ) + + def forward(self, x): + weight = torch.ops.quantized_decomposed.dequantize_per_tensor( + self.weight, 0.02, 0, -128, 127, torch.int8 + ) + return torch.nn.functional.conv_transpose1d(x, weight, stride=2) + + self._check_left_alone(QuantizedDeconv(), (torch.randn(1, 4, 6),)) diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index 7f471832a30..be74338d754 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -57,9 +57,7 @@ cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --s ``` XNNPACK delegates the linear and convolution ops; the flow LM's attention runs -as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. The export also -rewrites the Mimi decoder's transposed convolutions as convolutions, which -XNNPACK runs faster. +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. `--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow LM's transformer and flow net with the scheme of the backend; the codec stays in diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index 2bae25d8f47..5794777bfb9 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -82,7 +82,7 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") return [st.use_webgpu_compat_ops] return { - "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d], + "xnnpack": [st.depthwise_conv_transpose_as_mul], "mlx": [st.decompose_elu], "vulkan": [ st.decompose_elu, diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py index 4c6a1dcdb9b..78efe551120 100644 --- a/examples/models/phonon/source_transformations.py +++ b/examples/models/phonon/source_transformations.py @@ -15,9 +15,9 @@ the Mimi decoder splits at every ELU. conv_transpose1d_as_2d Vulkan runs transposed convolution on 4-D inputs only. conv_transpose1d_as_conv1d - XNNPACK's deconvolution is slow; this runs the Mimi - decoder's transposed convolutions (kernel = - 2 * stride) as a convolution and an interleave. + Runs the Mimi decoder's transposed convolutions + (kernel = 2 * stride) as a convolution and an + interleave; XNNPACK's lowering does the same. depthwise_conv_transpose_as_mul Vulkan computes the Mimi upsampler (depthwise, kernel = 2 * stride) incorrectly; this runs it as a diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index 63b79dee939..c10d2201352 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -255,8 +255,9 @@ def test_vulkan_swaps_preserve_outputs(self): self._check_swaps("vulkan", [7, 1, 2]) def test_xnnpack_swaps_preserve_outputs(self): - # The depthwise upsampler, 2 decoder transposed convs. - self._check_swaps("xnnpack", [1, 2]) + # The depthwise upsampler; XNNPACK lowers the decoder's transposed + # convs as convs itself. + self._check_swaps("xnnpack", [1]) def test_conv_transpose_as_conv_matches(self): torch.manual_seed(0)