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 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..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 b5267d2405a..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,6 +52,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() @@ -168,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") @@ -176,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, diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 056c8b82a85..32be6ae872d 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); }; /* @@ -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/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)