From aebfa9b20cf891a75850cec84a6ca12c3d9e8247 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 11:47:37 -0800 Subject: [PATCH 1/3] XNNPACK: fix enable_bf16 load failures for bmm and ops without bf16 support XnnpackPartitioner(enable_bf16=True) lets XNNPACK run bf16 nodes, so bf16 linear layers stop running on the single-threaded portable kernels. But a program lowered this way can fail when it loads: 1. The partitioner hands XNNPACK bf16 nodes that XNNPACK cannot define in bf16: softmax, mean, average pooling, max pooling, bilinear upsampling and batch norm. 2. Batch matrix multiply is defined with a bf16 output, but XNNPACK only multiplies bf16 matrices into an fp32 result. A partitioner config can now say it does not support bf16 (supports_bf16 = False), and enable_bf16 adds bf16 only for the configs that support it. The configs of the ops above set it, and convolution moves its own bf16 check to the same flag. At runtime a bf16 batch matrix multiply writes an fp32 result and converts it to the bf16 output, the pattern the linear layer already used, now shared by both ops. --- .../xnnpack/partition/config/gemm_configs.py | 7 +- .../partition/config/generic_node_configs.py | 5 + .../xnnpack/partition/config/node_configs.py | 1 + .../partition/config/xnnpack_config.py | 7 +- backends/xnnpack/runtime/XNNCompiler.cpp | 149 +++++++++++------- .../xnnpack/test/test_xnnpack_partitioner.py | 81 ++++++++++ 6 files changed, 187 insertions(+), 63 deletions(-) diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index 1796b5bdbbe..018df9e3ddc 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -342,6 +342,7 @@ def supported_precision_types(self): class ConvolutionConfig(GEMMConfig): target_name = "convolution.default" + supports_bf16 = False def __init__(self, **kwargs): super().__init__( @@ -361,12 +362,6 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: kernel_node = get_input_node(node, 1) kernel_shape = get_shape(kernel_node) - if any( - value.meta["val"].dtype == torch.bfloat16 - for value in (node, kernel_node, get_input_node(node, 0)) - ): - why(node, "XNNPACK does not support BF16 convolution") - return False # The weight rank is the only reliable indicator of the conv dimensionality. # stride, padding and dilation may each be a single value that ATen # broadcasts over every spatial dim. diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index a9a323d9453..934146cc20e 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -234,6 +234,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: class AvgPoolingConfig(GenericNodePartitionerConfig): target_name = "avg_pool2d.default" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ @@ -385,6 +386,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: class SoftmaxConfig(GenericNodePartitionerConfig): target_name = "_softmax.default" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ @@ -462,6 +464,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: class MaxPool2dConfig(GenericNodePartitionerConfig): target_name = "max_pool2d.default" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ @@ -496,6 +499,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]: class UpsampleBilinear2dConfig(GenericNodePartitionerConfig): target_name = "upsample_bilinear2d.vec" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ @@ -678,6 +682,7 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: class MeanDimConfig(GenericNodePartitionerConfig): target_name = "mean.dim" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ diff --git a/backends/xnnpack/partition/config/node_configs.py b/backends/xnnpack/partition/config/node_configs.py index bffe87d30c5..669a2578e0b 100644 --- a/backends/xnnpack/partition/config/node_configs.py +++ b/backends/xnnpack/partition/config/node_configs.py @@ -25,6 +25,7 @@ class BatchNormConfig(XNNPartitionerConfig): target_name = "_native_batch_norm_legit_no_training.default" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: if not self.check_common_constraints(node, ep): diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index b5267d2405a..bb605b97ad5 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -45,6 +45,11 @@ class XNNPartitionerConfig(PartitionerConfig): types they want to enable """ + # XNNPACK defines some ops only for fp32 and fp16 tensors. Their configs set this + # to False, so enable_bf16 leaves their bf16 nodes to the portable kernels + # instead of handing XNNPACK a node it refuses when the program loads. + supports_bf16 = True + def __init__(self, **kwargs): super().__init__() self.enabled_precision_types = self.supported_precision_types() @@ -269,7 +274,7 @@ def _check_node_has_valid_dtype(self, node): torch.float32, torch.float16, } - if self.enable_bf16: + if self.enable_bf16 and self.supports_bf16: valid_dtypes.add(torch.bfloat16) # Only allow int8 and quant dtypes for quant operations if is_quant(node) or is_dequant(node) or is_qparam(node): diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 056c8b82a85..b2ddcb98df4 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -818,6 +818,74 @@ const fb_xnnpack::XNNTensorValue* getSerializedTensorValue( return nullptr; } +/* +XNNPACK multiplies bf16 matrices only into an fp32 result. When the serialized +graph multiplies bf16 matrices into a bf16 output, define an fp32 value for the +op to write instead, which convertBf16MatMulResult() then converts into the +output, so the delegate boundary stays bf16. Otherwise the op writes the output +itself. +*/ +Error defineBf16MatMulResult( + xnn_subgraph_t subgraph_ptr, + const fb_xnnpack::XNNGraph* graph, + uint32_t input1_id, + uint32_t input2_id, + uint32_t output_id, + uint32_t remapped_output_id, + const NodePtr node, + uint32_t* result_id) noexcept { + *result_id = remapped_output_id; + const auto* input1 = getSerializedTensorValue(graph, input1_id); + const auto* input2 = getSerializedTensorValue(graph, input2_id); + const auto* output = getSerializedTensorValue(graph, output_id); + if (input1 == nullptr || input2 == nullptr || output == nullptr || + input1->datatype() != DataType::xnn_datatype_bf16 || + input2->datatype() != DataType::xnn_datatype_bf16 || + output->datatype() != DataType::xnn_datatype_bf16) { + return Error::Ok; + } + std::vector dims = flatbufferDimsToVector(output->dims()); + xnn_status status = xnn_define_tensor_value( + subgraph_ptr, + xnn_datatype_fp32, + dims.size(), + dims.data(), + /*data=*/nullptr, + /*external_id=*/XNN_INVALID_VALUE_ID, + /*flags=*/0, + result_id); + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Failed to define the fp32 result of bf16 node %i: %s", + node->debug_handle(), + xnn_status_to_string(status)); + return Error::Ok; +} + +/* +Converts the fp32 result that defineBf16MatMulResult() defined into the bf16 +output. Does nothing when the op wrote the output itself. +*/ +Error convertBf16MatMulResult( + xnn_subgraph_t subgraph_ptr, + uint32_t result_id, + uint32_t output_id, + const NodePtr node) noexcept { + if (result_id == output_id) { + return Error::Ok; + } + xnn_status status = + xnn_define_convert(subgraph_ptr, result_id, output_id, /*flags=*/0); + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Failed to convert the result of bf16 node %i to bf16: %s", + node->debug_handle(), + xnn_status_to_string(status)); + return Error::Ok; +} + /* Define serialized linear(fully-connected) node into the subgraph using the remapped ids to map the serialized ids, to the new ids generated @@ -828,8 +896,6 @@ Error defineFullyConnectedNode( 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_XNNFullyConnected(); std::pair min_max = getOutputMinMax(node); REMAP_ID(remapped_ids, graph_node->input1_id(), fc_input1); @@ -837,43 +903,16 @@ Error defineFullyConnectedNode( REMAP_ID(remapped_ids, graph_node->bias_id(), fc_bias); REMAP_ID(remapped_ids, graph_node->output_id(), fc_output); - // XNNPACK only provides a bf16 fully-connected of type bf16_bf16_f32: - // bf16 activation x bf16 weight -> fp32 output. When the serialized graph - // asks for a bf16 output (e.g. a fully bf16 model), define the FC with an - // fp32 intermediate output and append a convert (fp32 -> bf16) so the - // delegate boundary stays bf16. - const auto* in_tv = getSerializedTensorValue(graph, graph_node->input1_id()); - const auto* filt_tv = - getSerializedTensorValue(graph, graph_node->filter_id()); - const auto* out_tv = getSerializedTensorValue(graph, graph_node->output_id()); - const bool needs_bf16_output_convert = in_tv != nullptr && - filt_tv != nullptr && out_tv != nullptr && - in_tv->datatype() == DataType::xnn_datatype_bf16 && - filt_tv->datatype() == DataType::xnn_datatype_bf16 && - out_tv->datatype() == DataType::xnn_datatype_bf16; - - uint32_t fc_compute_output = fc_output; - if (needs_bf16_output_convert) { - std::vector out_dims = - flatbufferDimsToVector(out_tv->dims()); - uint32_t intermediate_id = XNN_INVALID_VALUE_ID; - xnn_status ts = xnn_define_tensor_value( - subgraph_ptr, - xnn_datatype_fp32, - out_dims.size(), - out_dims.data(), - /*data=*/nullptr, - /*external_id=*/XNN_INVALID_VALUE_ID, - /*flags=*/0, - &intermediate_id); - ET_CHECK_OR_RETURN_ERROR( - ts == xnn_status_success, - Internal, - "Failed to define fp32 intermediate for bf16 linear node %i: %s", - node->debug_handle(), - xnn_status_to_string(ts)); - fc_compute_output = intermediate_id; - } + uint32_t fc_result = XNN_INVALID_VALUE_ID; + ET_CHECK_OK_OR_RETURN_ERROR(defineBf16MatMulResult( + subgraph_ptr, + graph, + graph_node->input1_id(), + graph_node->filter_id(), + graph_node->output_id(), + fc_output, + node, + &fc_result)); xnn_status status = xnn_define_fully_connected( subgraph_ptr, @@ -882,7 +921,7 @@ Error defineFullyConnectedNode( fc_input1, fc_filter, fc_bias, - fc_compute_output, + fc_result, graph_node->flags()); ET_CHECK_OR_RETURN_ERROR( status == xnn_status_success, @@ -891,18 +930,7 @@ Error defineFullyConnectedNode( node->debug_handle(), xnn_status_to_string(status)); - if (needs_bf16_output_convert) { - xnn_status cs = xnn_define_convert( - subgraph_ptr, fc_compute_output, fc_output, /*flags=*/0); - ET_CHECK_OR_RETURN_ERROR( - cs == xnn_status_success, - Internal, - "Failed to define bf16 output convert for linear node %i: %s", - node->debug_handle(), - xnn_status_to_string(cs)); - } - - return Error::Ok; + return convertBf16MatMulResult(subgraph_ptr, fc_result, fc_output, node); }; /* @@ -1735,15 +1763,24 @@ Error defineBatchMatrixMultiplyNode( 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_XNNBatchMatrixMultiply(); REMAP_ID(remapped_ids, graph_node->input1_id(), bmm_in1); REMAP_ID(remapped_ids, graph_node->input2_id(), bmm_in2); REMAP_ID(remapped_ids, graph_node->output_id(), bmm_out); + uint32_t bmm_result = XNN_INVALID_VALUE_ID; + ET_CHECK_OK_OR_RETURN_ERROR(defineBf16MatMulResult( + subgraph_ptr, + graph, + graph_node->input1_id(), + graph_node->input2_id(), + graph_node->output_id(), + bmm_out, + node, + &bmm_result)); + xnn_status status = xnn_define_batch_matrix_multiply( - subgraph_ptr, bmm_in1, bmm_in2, bmm_out, graph_node->flags()); + subgraph_ptr, bmm_in1, bmm_in2, bmm_result, graph_node->flags()); ET_CHECK_OR_RETURN_ERROR( status == xnn_status_success, @@ -1752,7 +1789,7 @@ Error defineBatchMatrixMultiplyNode( node->debug_handle(), xnn_status_to_string(status)); - return Error::Ok; + return convertBf16MatMulResult(subgraph_ptr, bmm_result, bmm_out, node); } /* diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..3559bbfa943 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -161,3 +161,84 @@ def example_inputs(self): fwd2_et = executorch_module.run_method("forward_2", example_inputs) self.assertTrue(torch.allclose(fwd1_eager, fwd1_et[0], 1e-3)) self.assertTrue(torch.allclose(fwd2_eager, fwd2_et[0], 1e-3)) + + def test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16(self): + """ + XNNPACK refuses these ops on bf16 tensors when the program loads, so even + with enable_bf16 their bf16 nodes must run on the portable kernels. + """ + + class Mean(torch.nn.Module): + def forward(self, x): + return x.mean((-1, -2), keepdim=True) + + class Softmax(torch.nn.Module): + def forward(self, x): + return torch.softmax(x, -1) + + modules = { + "avg_pool2d": torch.nn.AvgPool2d(2), + "batch_norm": torch.nn.BatchNorm2d(4), + "max_pool2d": torch.nn.MaxPool2d(2), + "mean": Mean(), + "softmax": Softmax(), + "upsample_bilinear2d": torch.nn.Upsample(scale_factor=2, mode="bilinear"), + } + x = torch.randn(1, 4, 8, 8, dtype=torch.bfloat16) + for name, module in modules.items(): + with self.subTest(name): + module = module.eval().to(torch.bfloat16) + program = to_edge_transform_and_lower( + export(module, (x,)), + partitioner=[XnnpackPartitioner(enable_bf16=True)], + ) + self.assertFalse( + any( + "lowered_module" in node.name + for node in program.exported_program().graph.nodes + ) + ) + output = _load_for_executorch_from_buffer( + program.to_executorch().buffer + ).forward((x,))[0] + torch.testing.assert_close(output, module(x)) + + def test_enable_bf16_runs_an_attention_block(self): + """ + The linear layers and both batch matrix multiplications of a bf16 attention + block run in XNNPACK, and the program loads and matches eager bf16. + """ + + class Attention(torch.nn.Module): + def __init__(self): + super().__init__() + self.q = torch.nn.Linear(16, 16) + self.k = torch.nn.Linear(16, 16) + self.v = torch.nn.Linear(16, 16) + self.o = torch.nn.Linear(16, 16, bias=False) + + def forward(self, x): + scores = self.q(x) @ self.k(x).transpose(1, 2) * 0.25 + return self.o(torch.softmax(scores, dim=-1) @ self.v(x)) + + module = Attention().eval().to(torch.bfloat16) + x = torch.randn(2, 8, 16, dtype=torch.bfloat16) + program = to_edge_transform_and_lower( + export(module, (x,)), partitioner=[XnnpackPartitioner(enable_bf16=True)] + ) + left_on_cpu = { + node.target.__name__ + for node in program.exported_program().graph.nodes + if node.op == "call_function" and hasattr(node.target, "__name__") + } + matmuls = { + "aten.addmm.default", + "aten.bmm.default", + "aten.linear.default", + "aten.mm.default", + } + self.assertEqual(left_on_cpu & matmuls, set()) + output = _load_for_executorch_from_buffer( + program.to_executorch().buffer + ).forward((x,))[0] + torch.testing.assert_close(output, module(x), atol=1e-2, rtol=1e-2) From e7fa599c49f232643498eb419dedb80d98b114ec Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 13:23:04 -0800 Subject: [PATCH 2/3] XNNPACK: let a bf16 BatchNorm fuse into a linear, and tighten the bf16 tests A bf16 BatchNorm that follows a linear folds into the linear's weights, and the result is a bf16 fully-connected, which XNNPACK runs. BatchNormConfig now keeps a bf16 BatchNorm when FuseBatchNormPass.can_fuse accepts the linear before it, instead of refusing every bf16 BatchNorm. With per_op_mode the BatchNorm still gets its own delegate, so it cannot fold and the program fails to load, as it does without this change. FuseBatchNormPass.can_fuse now also checks that the BatchNorm is the only user of its input, as the pass itself requires before it fuses. Before, can_fuse accepted a linear whose output had other users, so the partitioner kept a bf16 BatchNorm that the pass then left unfused, and the program failed to load. A config that sets supports_bf16 = False now logs why it refuses a bf16 node, as the convolution config did before, and the comment notes that the list follows the pinned XNNPACK commit. convertBf16MatMulResult() now names its output parameter remapped_output_id, since it takes the remapped id. Tests: - test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16 also lowers each module in fp32 and checks that it is delegated, so the test shows that bf16 is what keeps it out. It adds a convolution followed by a BatchNorm, which must stay out too. - test_enable_bf16_fuses_batch_norm_into_a_linear checks that a bf16 linear followed by BatchNorm1d runs entirely in XNNPACK, and that the BatchNorm stays out when the linear output is also used elsewhere. - test_enable_bf16_runs_an_attention_block records the number of delegates and checks that the softmax stays out. - test_bmm.py gets a bf16 case, like the bf16 cases in the other op tests. --- backends/xnnpack/_passes/fuse_batch_norm.py | 6 ++ .../xnnpack/partition/config/node_configs.py | 13 ++- .../partition/config/xnnpack_config.py | 27 ++++-- backends/xnnpack/runtime/XNNCompiler.cpp | 8 +- backends/xnnpack/test/ops/test_bmm.py | 23 ++++- .../xnnpack/test/test_xnnpack_partitioner.py | 87 +++++++++++++++---- 6 files changed, 134 insertions(+), 30 deletions(-) diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py index 76ecce91585..2344d4bb835 100644 --- a/backends/xnnpack/_passes/fuse_batch_norm.py +++ b/backends/xnnpack/_passes/fuse_batch_norm.py @@ -102,6 +102,12 @@ def can_fuse( # noqa: C901 is_conv = input_node.target == exir_ops.edge.aten.convolution.default + # Fusing changes what the input node computes, so nothing else may use it. + if len(input_node.users) != 1: + if why: + why(bn, "Batch norm must be the only user of the input node.") + return False + # All users of the batch_norm node must be getitem ops. # batch_norm returns a 3-element tuple. # Each user must only access the first element of the tuple. diff --git a/backends/xnnpack/partition/config/node_configs.py b/backends/xnnpack/partition/config/node_configs.py index 669a2578e0b..782a336053b 100644 --- a/backends/xnnpack/partition/config/node_configs.py +++ b/backends/xnnpack/partition/config/node_configs.py @@ -17,6 +17,7 @@ ) from executorch.backends.xnnpack.utils.utils import is_param_node from executorch.exir.backend.utils import WhyNoPartition +from executorch.exir.dialects._ops import ops as exir_ops from torch.export import ExportedProgram logger = logging.getLogger(__name__) @@ -25,7 +26,6 @@ class BatchNormConfig(XNNPartitionerConfig): target_name = "_native_batch_norm_legit_no_training.default" - supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: if not self.check_common_constraints(node, ep): @@ -41,6 +41,17 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: why(node, f"BatchNorm cannot be decomposed or fused with {input_node}") return False + # Unless it fuses into a linear, a BatchNorm becomes a depthwise convolution + # or folds into a convolution, and the pinned XNNPACK defines neither in bf16. + if input_node.meta["val"].dtype == torch.bfloat16 and not ( + can_fuse and input_node.target == exir_ops.edge.aten.linear.default + ): + why( + node, + "XNNPACK supports a bf16 BatchNorm only when it fuses into a linear", + ) + return False + return True def get_node_and_deps( diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index bb605b97ad5..e1bc5a4ab0a 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -35,6 +35,13 @@ class ConfigPrecisionType(Enum): DYNAMIC_QUANT = 3 +def _input_and_output_values(node: torch.fx.Node) -> List[object]: + values = [arg.meta.get("val") for arg in node.all_input_nodes] + node_val = node.meta.get("val") + values.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val]) + return values + + class XNNPartitionerConfig(PartitionerConfig): """ Base partitioner config for XNNPACK Partitioner Configs. Base wrapper class @@ -45,9 +52,9 @@ class XNNPartitionerConfig(PartitionerConfig): types they want to enable """ - # XNNPACK defines some ops only for fp32 and fp16 tensors. Their configs set this - # to False, so enable_bf16 leaves their bf16 nodes to the portable kernels - # instead of handing XNNPACK a node it refuses when the program loads. + # At the pinned XNNPACK commit, some ops are defined only for fp32 and fp16 + # tensors. Their configs set this to False, so their bf16 nodes stay on the + # portable kernels instead of failing when the program loads. supports_bf16 = True def __init__(self, **kwargs): @@ -173,6 +180,13 @@ def check_common_constraints( why(node, reason="not enabled precision types") return False + if not self.supports_bf16 and any( + isinstance(val, torch.Tensor) and val.dtype == torch.bfloat16 + for val in _input_and_output_values(node) + ): + why(node, reason="XNNPACK does not support bf16 for this op") + return False + has_valid_dtypes = self._check_node_has_valid_dtype(node) if not has_valid_dtypes: why(node, reason="invalid dtype") @@ -181,10 +195,7 @@ def check_common_constraints( return self._check_tensor_ranks(node) def _check_tensor_ranks(self, node: torch.fx.Node) -> bool: - vals = [arg.meta.get("val") for arg in node.all_input_nodes] - node_val = node.meta.get("val") - vals.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val]) - for val in vals: + for val in _input_and_output_values(node): if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS: why( node, @@ -274,7 +285,7 @@ def _check_node_has_valid_dtype(self, node): torch.float32, torch.float16, } - if self.enable_bf16 and self.supports_bf16: + if self.enable_bf16: valid_dtypes.add(torch.bfloat16) # Only allow int8 and quant dtypes for quant operations if is_quant(node) or is_dequant(node) or is_qparam(node): diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index b2ddcb98df4..32be6ae872d 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -870,13 +870,13 @@ output. Does nothing when the op wrote the output itself. Error convertBf16MatMulResult( xnn_subgraph_t subgraph_ptr, uint32_t result_id, - uint32_t output_id, + uint32_t remapped_output_id, const NodePtr node) noexcept { - if (result_id == output_id) { + if (result_id == remapped_output_id) { return Error::Ok; } - xnn_status status = - xnn_define_convert(subgraph_ptr, result_id, output_id, /*flags=*/0); + xnn_status status = xnn_define_convert( + subgraph_ptr, result_id, remapped_output_id, /*flags=*/0); ET_CHECK_OR_RETURN_ERROR( status == xnn_status_success, Internal, diff --git a/backends/xnnpack/test/ops/test_bmm.py b/backends/xnnpack/test/ops/test_bmm.py index a029738e771..49d0112a19c 100644 --- a/backends/xnnpack/test/ops/test_bmm.py +++ b/backends/xnnpack/test/ops/test_bmm.py @@ -7,7 +7,8 @@ import unittest import torch -from executorch.backends.xnnpack.test.tester import Tester +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.backends.xnnpack.test.tester import Tester, ToEdgeTransformAndLower class TestBMM(unittest.TestCase): @@ -47,3 +48,23 @@ def test_fp32_bmm(self): torch.randn(2, 4, 6), ) self._test_bmm(inputs) + + def test_bf16_bmm(self): + inputs = ( + torch.randn(2, 3, 4).to(torch.bfloat16), + torch.randn(2, 4, 6).to(torch.bfloat16), + ) + ( + Tester(self.BMM(), inputs) + .export() + .to_edge_transform_and_lower( + ToEdgeTransformAndLower( + partitioners=[XnnpackPartitioner(enable_bf16=True)] + ) + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not(["executorch_exir_dialects_edge__ops_aten_bmm_default"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(atol=0.01, rtol=0.01) + ) diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 3559bbfa943..be775e5d784 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -19,6 +19,14 @@ from torch.export import export +def _ops_outside_delegates(program): + return [ + node.target.__name__ + for node in program.exported_program().graph.nodes + if node.op == "call_function" and hasattr(node.target, "__name__") + ] + + class TestXnnpackPartitioner(unittest.TestCase): """Test cases for XnnpackPartitioner functionality and deprecation warnings.""" @@ -165,7 +173,8 @@ def example_inputs(self): def test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16(self): """ XNNPACK refuses these ops on bf16 tensors when the program loads, so even - with enable_bf16 their bf16 nodes must run on the portable kernels. + with enable_bf16 their bf16 nodes must run on the portable kernels. The + same modules are delegated in fp32, so bf16 is what keeps them out. """ class Mean(torch.nn.Module): @@ -176,32 +185,79 @@ class Softmax(torch.nn.Module): def forward(self, x): return torch.softmax(x, -1) + def lower(module, x): + return to_edge_transform_and_lower( + export(module, (x,)), + partitioner=[XnnpackPartitioner(enable_bf16=True)], + ) + modules = { "avg_pool2d": torch.nn.AvgPool2d(2), "batch_norm": torch.nn.BatchNorm2d(4), + "conv_batch_norm": torch.nn.Sequential( + torch.nn.Conv2d(4, 4, 3), torch.nn.BatchNorm2d(4) + ), "max_pool2d": torch.nn.MaxPool2d(2), "mean": Mean(), "softmax": Softmax(), "upsample_bilinear2d": torch.nn.Upsample(scale_factor=2, mode="bilinear"), } - x = torch.randn(1, 4, 8, 8, dtype=torch.bfloat16) + x = torch.randn(1, 4, 8, 8) for name, module in modules.items(): with self.subTest(name): - module = module.eval().to(torch.bfloat16) + module.eval() + self.assertIn( + "executorch_call_delegate", _ops_outside_delegates(lower(module, x)) + ) + module, bf16_x = module.to(torch.bfloat16), x.to(torch.bfloat16) + program = lower(module, bf16_x) + self.assertNotIn( + "executorch_call_delegate", _ops_outside_delegates(program) + ) + output = _load_for_executorch_from_buffer( + program.to_executorch().buffer + ).forward((bf16_x,))[0] + torch.testing.assert_close(output, module(bf16_x)) + + def test_enable_bf16_fuses_batch_norm_into_a_linear(self): + """ + XNNPACK cannot define a bf16 BatchNorm on its own, but one that follows a + linear folds into it, so the whole bf16 block still runs in XNNPACK. It + does not fold when the linear output has another user, so then it stays + on the portable kernels. + """ + + class LinearBatchNorm(torch.nn.Module): + def __init__(self, reuse_linear_output): + super().__init__() + self.linear = torch.nn.Linear(8, 4) + self.batch_norm = torch.nn.BatchNorm1d(4) + self.batch_norm.running_mean.uniform_(-1, 1) + self.batch_norm.running_var.uniform_(0.5, 2) + self.reuse_linear_output = reuse_linear_output + + def forward(self, x): + y = self.linear(x) + z = self.batch_norm(y) + return z + y if self.reuse_linear_output else z + + x = torch.randn(2, 8, dtype=torch.bfloat16) + delegated = {"executorch_call_delegate", "getitem"} + for reuse_linear_output, ops_outside in ( + (False, delegated), + (True, delegated | {"aten._native_batch_norm_legit_no_training.default"}), + ): + with self.subTest(reuse_linear_output=reuse_linear_output): + module = LinearBatchNorm(reuse_linear_output).eval().to(torch.bfloat16) program = to_edge_transform_and_lower( export(module, (x,)), partitioner=[XnnpackPartitioner(enable_bf16=True)], ) - self.assertFalse( - any( - "lowered_module" in node.name - for node in program.exported_program().graph.nodes - ) - ) + self.assertEqual(set(_ops_outside_delegates(program)), ops_outside) output = _load_for_executorch_from_buffer( program.to_executorch().buffer ).forward((x,))[0] - torch.testing.assert_close(output, module(x)) + torch.testing.assert_close(output, module(x), atol=1e-2, rtol=1e-2) def test_enable_bf16_runs_an_attention_block(self): """ @@ -226,18 +282,17 @@ def forward(self, x): program = to_edge_transform_and_lower( export(module, (x,)), partitioner=[XnnpackPartitioner(enable_bf16=True)] ) - left_on_cpu = { - node.target.__name__ - for node in program.exported_program().graph.nodes - if node.op == "call_function" and hasattr(node.target, "__name__") - } + ops = _ops_outside_delegates(program) + # The expand copies that matmul adds split the block into 3 delegates. + self.assertEqual(ops.count("executorch_call_delegate"), 3) + self.assertIn("aten._softmax.default", ops) matmuls = { "aten.addmm.default", "aten.bmm.default", "aten.linear.default", "aten.mm.default", } - self.assertEqual(left_on_cpu & matmuls, set()) + self.assertEqual(set(ops) & matmuls, set()) output = _load_for_executorch_from_buffer( program.to_executorch().buffer ).forward((x,))[0] From fdbbcb598e4b8e86f7495bd1be1e2906ee5ada87 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 14:30:33 -0800 Subject: [PATCH 3/3] XNNPACK: fold a BatchNorm into a linear only when the linear output is 2-D BatchNorm normalizes dim 1. For a 2-D linear output, dim 1 holds the linear's output features, so FuseBatchNormPass can fold the BatchNorm into the linear's weights. For a 3-D output it does not. can_fuse still accepted that case, so when the sizes happened to match, the folded linear gave wrong results without any error, in fp32 and in bf16. When the sizes did not match, lowering failed inside the pass. can_fuse now accepts a linear only when its output is 2-D. Otherwise the BatchNorm is decomposed into a depthwise convolution, which normalizes dim 1, or in bf16 it stays on the portable kernels. Test: test_fp32_linear_batch_norm_on_3d_output lowers a Linear followed by BatchNorm1d on a 3-D input to XNNPACK, checks that the BatchNorm runs inside the delegate, runs the program and compares with eager. Without this change the outputs do not match. --- backends/xnnpack/_passes/fuse_batch_norm.py | 6 ++++++ .../test/passes/test_batch_norm_fusion.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py index 2344d4bb835..9e383a77c5b 100644 --- a/backends/xnnpack/_passes/fuse_batch_norm.py +++ b/backends/xnnpack/_passes/fuse_batch_norm.py @@ -147,6 +147,12 @@ def can_fuse( # noqa: C901 if why: why("Convolution input must be rank 3 or 4.") return False + # BatchNorm normalizes dim 1, which holds the linear's output features only + # when the linear output is 2-D. + elif "val" not in input_node.meta or len(input_node.meta["val"].shape) != 2: + if why: + why(bn, "Linear output must be rank 2.") + return False return True diff --git a/backends/xnnpack/test/passes/test_batch_norm_fusion.py b/backends/xnnpack/test/passes/test_batch_norm_fusion.py index 87c7fa397ef..519cf94243a 100644 --- a/backends/xnnpack/test/passes/test_batch_norm_fusion.py +++ b/backends/xnnpack/test/passes/test_batch_norm_fusion.py @@ -135,6 +135,24 @@ def test_fp32_linear_batch_norm_fusion(self): .run_method_and_compare_outputs() ) + def test_fp32_linear_batch_norm_on_3d_output(self): + """ + BatchNorm1d normalizes dim 1. When the linear output is 3-D, dim 1 is not the + linear's output features, so the batch norm must not fold into the linear. + """ + model = torch.nn.Sequential(torch.nn.Linear(8, 4), torch.nn.BatchNorm1d(4)) + model[1].running_mean.uniform_(-1, 1) + model[1].running_var.uniform_(0.5, 2) + ( + Tester(model.eval(), (torch.randn(2, 4, 8),)) + .export() + .to_edge_transform_and_lower() + .check_count({self.bn_name: 0}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_fp32_conv3d_batch_norm_doesnt_partition(self): """ Conv3d is not currently supported by XNNPACK. We also don't support standalone