[xnnpack] Support bf16 delegation for fully-connected

backends/xnnpack/cmake/Dependencies.cmake

@@ -66,6 +66,10 @@ if(WIN32)
   set_overridable_option(XNNPACK_ENABLE_AVX256VNNI OFF)
   set_overridable_option(XNNPACK_ENABLE_AVX256VNNIGFNI OFF)
   set_overridable_option(XNNPACK_ENABLE_AVX512BF16 OFF)
+  # clang-cl (reported by CMake as Clang, not MSVC, so XNNPACK's own MSVC gate
+  # does not catch it) crashes with an internal codegen error compiling some of
+  # the AVX512-FP16 micro-kernels, so disable them on Windows.
+  set_overridable_option(XNNPACK_ENABLE_AVX512FP16 OFF)
 endif()
 
 set(XNNPACK_BUILD_ALL_MICROKERNELS

backends/xnnpack/operators/node_visitor.py

@@ -207,7 +207,7 @@
 
         return ext_id, id_out, flag
 
     def get_serialized_dtype(
         self,
         quant_params: Optional[QuantParams],
         node: torch.fx.Node,
@@ -271,6 +271,12 @@
                     if force_fp32
                     else XNNDatatype.xnn_datatype_fp16
                 )
+            elif node_dtype is not None and node_dtype == torch.bfloat16:
+                dtype = (
+                    XNNDatatype.xnn_datatype_fp32
+                    if force_fp32
+                    else XNNDatatype.xnn_datatype_bf16
+                )
 
         return dtype
 
@@ -591,7 +597,7 @@
         # Quantize buffer if static data is indeed quantized
         if quant_params is not None and not quant_params.is_dynamic:
             const_val = quant_params.quantize_tensor(const_val).contiguous()
-        elif const_val.dtype != torch.float16 or force_fp32:
+        elif const_val.dtype not in (torch.float16, torch.bfloat16) or force_fp32:
             # ensure that the const is fp32
             const_val = const_val.to(dtype=torch.float32).contiguous()
 
@@ -712,12 +718,19 @@
                 bias_quant_params = QuantParams.from_bias(
                     bias_node, weight_quant_params, input_quant_params
                 )
+                # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16
+                # activation/weight but an fp32 bias, so force the bias to fp32.
+                weight_val = weight_node.meta.get("val", None)
+                bias_force_fp32 = (
+                    weight_val is not None and weight_val.dtype == torch.bfloat16
+                )
                 self.define_tensor(
                     bias_node,
                     xnn_graph,
                     vals_to_ids,
                     quant_params=bias_quant_params,
                     convert_to_nhwc=False,  # Bias is generally 1d and can not be in NHWC
+                    force_fp32=bias_force_fp32,
                 )
 
     def define_node(

backends/xnnpack/operators/op_linear.py

@@ -73,6 +73,11 @@ def define_node(
             force_fp32 = False
             if input_quant_params is not None and input_quant_params.is_dynamic:
                 force_fp32 = True
+            # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16
+            # activation/weight but an fp32 bias, so force the bias to fp32.
+            weight_val = weight_node.meta.get("val", None)
+            if weight_val is not None and weight_val.dtype == torch.bfloat16:
+                force_fp32 = True
 
             self.define_tensor(
                 get_input_node(node, 2),

backends/xnnpack/partition/config/xnnpack_config.py

@@ -228,6 +228,7 @@ def _check_node_has_valid_dtype(self, node):
         valid_dtypes = {
             torch.float32,
             torch.float16,
+            torch.bfloat16,
         }
         # Only allow int8 and quant dtypes for quant operations
         if is_quant(node) or is_dequant(node) or is_qparam(node):

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -791,6 +791,33 @@
 
   return Error::Ok;
 };
+/*
+Look up a serialized tensor value (plain or quantized wrapper) by its
+output id. Returns nullptr if not found.
+*/
+const fb_xnnpack::XNNTensorValue* getSerializedTensorValue(
+    const fb_xnnpack::XNNGraph* graph,
+    uint32_t id) noexcept {
+  if (graph == nullptr || graph->xvalues() == nullptr) {
+    return nullptr;
+  }
+  for (auto value : *graph->xvalues()) {
+    const fb_xnnpack::XNNTensorValue* tv = nullptr;
+    if (value->xvalue_union_type() ==
+        fb_xnnpack::XValueUnion::XNNTensorValue) {
+      tv = value->xvalue_union_as_XNNTensorValue();
+    } else if (
+        value->xvalue_union_type() ==
+        fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
+      tv = value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value();
+    }
+    if (tv != nullptr && tv->id_out() == id) {
+      return tv;
+    }
+  }
+  return nullptr;
+}
+
 /*
 Define serialized linear(fully-connected) node into the subgraph using
 the remapped ids to map the serialized ids, to the new ids generated
@@ -810,14 +837,52 @@
   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;
+  }
+
   xnn_status status = xnn_define_fully_connected(
       subgraph_ptr,
       min_max.first,
       min_max.second,
       fc_input1,
       fc_filter,
       fc_bias,
-      fc_output,
+      fc_compute_output,
       graph_node->flags());
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
@@ -826,6 +891,17 @@
       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;
 };
 

backends/xnnpack/third-party/xnnpack.buck.bzl

@@ -70,7 +70,10 @@ def define_xnnpack():
     # @lint-ignore BUCKLINT: native and fb_native are explicitly forbidden in fbcode.
     native.cxx_library(
         name = "subgraph",
-        srcs = SUBGRAPH_SRCS + ["XNNPACK/src/datatype.c"],
+        srcs = SUBGRAPH_SRCS + [
+            "XNNPACK/src/datatype.c",
+            "XNNPACK/src/subgraph/rewrites/cvt_to_fp32.cc",
+        ],
         compiler_flags = [
             "-Wno-error=missing-braces",  # required since the SGX toolchain does not have this by default
         ],
@@ -1150,6 +1153,7 @@ def define_xnnpack():
             "XNNPACK/src/init.c",
             "XNNPACK/src/params.c",
             "XNNPACK/src/configs/hardware-config.c",
+            "XNNPACK/src/xnnpack/init-once.c",
             "XNNPACK/src/microparams-init.c",
             "XNNPACK/src/microkernel-utils.c",
             "XNNPACK/src/reference/binary-elementwise.cc",