From 958ed92f66049cfe20d94afdcea78400be4d64cb Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 16:11:51 -0800 Subject: [PATCH] XNNPACK: fix enable_bf16 load failures and BatchNorm folding into a 3-D linear output (#24464) ### What is wrong today `XnnpackPartitioner(enable_bf16=True)` lets XNNPACK run bf16 nodes. Without it, bf16 matrix multiplies run on the single-threaded portable kernels, which is very slow. With it, a program can fail to load: 1. The partitioner hands XNNPACK bf16 nodes that XNNPACK cannot define in bf16: softmax, mean, average and max pooling, bilinear upsampling, and batch norm that does not fold into a linear. Loading fails with `xnn_status_invalid_parameter`. 2. Batch matrix multiply (`bmm`) is defined with a bf16 output, but XNNPACK only multiplies bf16 matrices into an fp32 result. Loading fails with `Failed to create BMM node`. A bf16 attention block written with `matmul` and `softmax` hits both. ### What this change does - A config can now say it does not support bf16 (`supports_bf16 = False`), so its bf16 nodes stay on the portable kernels and the partitioner logs why. Softmax, mean, both poolings, bilinear upsampling and convolution set it. The default is `True`: XNNPACK rewrites most bf16 ops to fp32 internally, and at the pinned commit only these reject bf16 when defined. A bf16 run of the other float ops found no other failure. - A bf16 batch norm is kept in XNNPACK only when `FuseBatchNormPass.can_fuse` says it folds into the linear before it. - `can_fuse` now also requires the batch norm to be the only user of the linear output, and the linear output to be 2-D. Without these checks, a bf16 batch norm could be kept but left unfused, so loading failed, and a batch norm after a 3-D linear output was folded along the wrong dim, which gave wrong results in fp32 and bf16. - A bf16 `bmm` now writes an fp32 result and converts it to bf16, as the linear layer already did. With this change, the SmolVLA robot policy (a vision-language-action model) in bf16 with `enable_bf16=True` loads, and then runs one action chunk in about 11 seconds on 16 x86 cores, instead of about 9 minutes with the default partitioner. ### Tests - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16`: each op, and a convolution followed by a batch norm, is delegated in fp32 but not in bf16, and the bf16 program matches eager. - `test_enable_bf16_fuses_batch_norm_into_a_linear`: a bf16 linear followed by `BatchNorm1d` runs fully in XNNPACK. When the linear output is also used elsewhere, the batch norm stays out. - `test_enable_bf16_runs_an_attention_block`: the linear layers and both `bmm`s run in XNNPACK, the softmax stays outside. - `test_bf16_bmm` in `test_bmm.py`, and `test_fp32_linear_batch_norm_on_3d_output` in `test_batch_norm_fusion.py`. All of these fail before this change, except the fused batch norm case, which guards the fusion that already worked. Ran the XNNPACK partitioner, op and pass tests on Linux x86_64. The only failures, in the per-channel int4 linear tests, also happen without this change. (cherry picked from commit 7bf0303039c6190b6081ac52811a7355609c9629) --- backends/xnnpack/_passes/fuse_batch_norm.py | 12 ++ .../xnnpack/partition/config/gemm_configs.py | 7 +- .../partition/config/generic_node_configs.py | 5 + .../xnnpack/partition/config/node_configs.py | 12 ++ .../partition/config/xnnpack_config.py | 30 ++++ backends/xnnpack/runtime/XNNCompiler.cpp | 149 +++++++++++------- backends/xnnpack/test/ops/test_bmm.py | 23 ++- .../test/passes/test_batch_norm_fusion.py | 18 +++ .../xnnpack/test/test_xnnpack_partitioner.py | 136 ++++++++++++++++ 9 files changed, 329 insertions(+), 63 deletions(-) diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py index 76ecce91585..9e383a77c5b 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. @@ -141,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/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index 4df6c4f1141..8afe3876c6e 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -345,6 +345,7 @@ def supported_precision_types(self): class ConvolutionConfig(GEMMConfig): target_name = "convolution.default" + supports_bf16 = False def __init__(self, **kwargs): super().__init__( @@ -364,12 +365,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 d445b5ca8ac..34898dfbe88 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -246,6 +246,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: """ @@ -372,6 +373,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]: class SoftmaxConfig(GenericNodePartitionerConfig): target_name = "_softmax.default" + supports_bf16 = False def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: """ @@ -449,6 +451,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: """ @@ -483,6 +486,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: """ @@ -636,6 +640,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..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__) @@ -40,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 e0a575e000a..44faace2af5 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -33,6 +33,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 @@ -43,6 +50,11 @@ class XNNPartitionerConfig(PartitionerConfig): types they want to enable """ + # 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): super().__init__() self.enabled_precision_types = self.supported_precision_types() @@ -147,11 +159,29 @@ 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") return False + return self._check_tensor_ranks(node) + + def _check_tensor_ranks(self, node: torch.fx.Node) -> bool: + for val in _input_and_output_values(node): + if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS: + why( + node, + reason=f"XNNPACK supports at most {XNN_MAX_TENSOR_DIMS} dims, got {val.dim()}", + ) + return False + return True def _check_inputs_are_valid_dtypes(self, node, valid_dtypes): diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 74a5eec4a91..8ae0dcb2e3b 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 remapped_output_id, + const NodePtr node) noexcept { + if (result_id == remapped_output_id) { + return Error::Ok; + } + 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, + "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); }; /* @@ -1702,15 +1730,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, @@ -1719,7 +1756,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/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/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 diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..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.""" @@ -161,3 +169,131 @@ 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. The + same modules are delegated in fp32, so bf16 is what keeps them out. + """ + + 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) + + 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) + for name, module in modules.items(): + with self.subTest(name): + 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.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), atol=1e-2, rtol=1e-2) + + 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)] + ) + 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(set(ops) & 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)