XNNPACK: fix enable_bf16 load failures and BatchNorm folding into a 3-D linear output

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
 

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.

backends/xnnpack/partition/config/generic_node_configs.py

@@ -233,6 +233,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:
         """
@@ -359,6 +360,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:
         """
@@ -436,6 +438,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:
         """
@@ -470,6 +473,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:
         """
@@ -623,6 +627,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:
         """

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(

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,

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<size_t> dims = flatbufferDimsToVector<size_t>(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,52 +896,23 @@ Error defineFullyConnectedNode(
     const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     const NodePtr node,
     const fb_xnnpack::XNNGraph* graph) noexcept {
-  MAYBE_UNUSED(graph);
-
   auto graph_node = node->xnode_union_as_XNNFullyConnected();
   std::pair<float, float> min_max = getOutputMinMax(node);
   REMAP_ID(remapped_ids, graph_node->input1_id(), fc_input1);
   REMAP_ID(remapped_ids, graph_node->filter_id(), fc_filter);
   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<size_t> out_dims =
-        flatbufferDimsToVector<size_t>(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<uint32_t, uint32_t>& 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);
 }
 
 /*