diff --git a/backends/xnnpack/cmake/Dependencies.cmake b/backends/xnnpack/cmake/Dependencies.cmake index 4fb80d5c18c..d02ac4a7926 100644 --- a/backends/xnnpack/cmake/Dependencies.cmake +++ b/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 diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index a0f03205ed5..7032b444e6a 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -271,6 +271,12 @@ def get_per_channel_dtype( 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 @@ def get_serialized_buffer_index( # 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 @@ def define_nodes_tensor_inputs_outputs( 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( diff --git a/backends/xnnpack/operators/op_linear.py b/backends/xnnpack/operators/op_linear.py index dda1d3e53ef..31eb0d7950f 100644 --- a/backends/xnnpack/operators/op_linear.py +++ b/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), diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index 817f9d1cf50..a9ad94ff438 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/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): diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 3f2fb5a1aaa..60b4ed6dcf4 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -791,6 +791,33 @@ Error defineConvertNode( 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,6 +837,44 @@ 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; + } + xnn_status status = xnn_define_fully_connected( subgraph_ptr, min_max.first, @@ -817,7 +882,7 @@ Error defineFullyConnectedNode( 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 @@ 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; }; diff --git a/backends/xnnpack/third-party/XNNPACK b/backends/xnnpack/third-party/XNNPACK index 1adaa7c709d..92a7ad501b9 160000 --- a/backends/xnnpack/third-party/XNNPACK +++ b/backends/xnnpack/third-party/XNNPACK @@ -1 +1 @@ -Subproject commit 1adaa7c709d4839d29e1f219cb962b01c9e6a905 +Subproject commit 92a7ad501b9516f9fecae119e0146dd1f58e54c1 diff --git a/backends/xnnpack/third-party/xnnpack.buck.bzl b/backends/xnnpack/third-party/xnnpack.buck.bzl index ac861435af8..3aa82f89cd4 100644 --- a/backends/xnnpack/third-party/xnnpack.buck.bzl +++ b/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",