[executorch][native-vulkan] Run bf16 methods in f32 on the Vulkan engine

backends/native/runtime/vulkan/VulkanEngineExecutable.cpp

@@ -51,6 +51,8 @@
 #include <executorch/backends/native/runtime/vulkan/passes/UseReferenceGgufKQuant.h>
 #include <executorch/runtime/core/freeable_buffer.h>
 
+#include <c10/util/BFloat16.h>
+
 #include <executorch/backends/vulkan/runtime/api/api.h>
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
@@ -116,8 +118,66 @@ vkapi::ScalarType to_vk_dtype(ScalarType t) {
       std::string("vulkan: unsupported dtype ") + scalar_type_name(t));
 }
 
+// ET-VK lacks int64 and bf16 storage, so int64 values run as int32 and bf16
+// values as f32 on the device.
 ScalarType device_dtype(ScalarType dtype) {
-  return dtype == ScalarType::Long ? ScalarType::Int : dtype;
+  if (dtype == ScalarType::Long) {
+    return ScalarType::Int;
+  }
+  if (dtype == ScalarType::BFloat16) {
+    return ScalarType::Float;
+  }
+  return dtype;
+}
+
+template <typename T>
+T load_element(const void* data, size_t i) {
+  T value;
+  std::memcpy(
+      &value, static_cast<const uint8_t*>(data) + i * sizeof(T), sizeof(T));
+  return value;
+}
+
+template <typename T>
+void store_element(void* data, size_t i, T value) {
+  std::memcpy(static_cast<uint8_t*>(data) + i * sizeof(T), &value, sizeof(T));
+}
+
+// Host copies into or out of bf16 go through f32 values rounded to bf16, so a
+// bf16 tensor holds bf16 values at the host boundary whatever the host dtype.
+void check_bf16_host_dtype(ScalarType dtype) {
+  if (dtype != ScalarType::BFloat16 && dtype != ScalarType::Float) {
+    throw std::runtime_error(
+        std::string("vulkan: bf16 copies need BFloat16 or Float host data, ") +
+        "not " + scalar_type_name(dtype));
+  }
+}
+
+std::vector<float>
+bf16_values_from_host(const void* data, size_t numel, ScalarType dtype) {
+  check_bf16_host_dtype(dtype);
+  std::vector<float> values(numel);
+  for (size_t i = 0; i < numel; ++i) {
+    values[i] = dtype == ScalarType::BFloat16
+        ? load_element<c10::BFloat16>(data, i)
+        : c10::BFloat16(load_element<float>(data, i));
+  }
+  return values;
+}
+
+void bf16_values_to_host(
+    const std::vector<float>& values,
+    void* data,
+    ScalarType dtype) {
+  check_bf16_host_dtype(dtype);
+  for (size_t i = 0; i < values.size(); ++i) {
+    const c10::BFloat16 value(values[i]);
+    if (dtype == ScalarType::BFloat16) {
+      store_element(data, i, value);
+    } else {
+      store_element(data, i, static_cast<float>(value));
+    }
+  }
 }
 
 size_t logical_numel(const TensorMeta& meta) {
@@ -649,7 +709,8 @@ class VulkanEngineExecutable final : public EngineExecutable {
   std::vector<uint8_t> zero_state_;
   // Bytes shared by tensor references until ET-VK consumes and releases them
   // during prepack: raw package bytes by canonical owner, and bytes computed
-  // from them (int32-narrowed Long constants, q4 repacks) by derivation key.
+  // from them (int32-narrowed Long constants, f32-widened bf16 constants, q4
+  // repacks) by derivation key.
   std::unordered_map<std::string, std::shared_ptr<OwnedBytes>>
       source_constants_;
   std::unordered_map<std::string, std::shared_ptr<OwnedBytes>>
@@ -792,8 +853,17 @@ class VulkanEngineExecutable final : public EngineExecutable {
       check_cache_position_input(data, numel, src_dtype, *bound, rows_written);
       pending_cache_rows_written_.at(i) = rows_written;
     }
+    const VkRef staging = input_staging_.at(i);
+    if (meta.dtype == ScalarType::BFloat16 ||
+        src_dtype == ScalarType::BFloat16) {
+      const std::vector<float> values =
+          bf16_values_from_host(data, numel, src_dtype);
+      graph_->maybe_cast_and_copy_into_staging(
+          staging, values.data(), numel, vkapi::kFloat);
+      return;
+    }
     graph_->maybe_cast_and_copy_into_staging(
-        input_staging_.at(i), data, numel, to_vk_dtype(src_dtype));
+        staging, data, numel, to_vk_dtype(src_dtype));
   }
 
   void execute() override {
@@ -815,11 +885,17 @@ class VulkanEngineExecutable final : public EngineExecutable {
     if (data == nullptr) {
       throw std::runtime_error("vulkan: output data is null");
     }
+    const VkRef staging = output_staging_.at(output_position);
+    if (output_dtype(i) == ScalarType::BFloat16 ||
+        dst_dtype == ScalarType::BFloat16) {
+      std::vector<float> values(numel);
+      graph_->maybe_cast_and_copy_from_staging(
+          staging, values.data(), numel, vkapi::kFloat);
+      bf16_values_to_host(values, data, dst_dtype);
+      return;
+    }
     graph_->maybe_cast_and_copy_from_staging(
-        output_staging_.at(output_position),
-        data,
-        numel,
-        to_vk_dtype(dst_dtype));
+        staging, data, numel, to_vk_dtype(dst_dtype));
   }
 
  private:
@@ -857,7 +933,8 @@ class VulkanEngineExecutable final : public EngineExecutable {
     }
     const auto align4 = [](int64_t v) { return (v + 3) / 4 * 4; };
     const int64_t limit = graph_->max_buffer_numel();
-    const int64_t elem = static_cast<int64_t>(element_size(query->dtype));
+    const int64_t elem =
+        static_cast<int64_t>(element_size(device_dtype(query->dtype)));
     int64_t weights = 0;
     if (__builtin_mul_overflow(
             query->sizes[2] * align4(query->sizes[1]),
@@ -1255,8 +1332,29 @@ class VulkanEngineExecutable final : public EngineExecutable {
             OwnedBytes::from_vector(std::move(converted)));
       }
       storage = narrowed;
+    } else if (info->dtype == ScalarType::BFloat16) {
+      std::shared_ptr<OwnedBytes>& widened = derived_constants_
+          [std::to_string(info->owner.size()) + ":" + info->owner + ":f32"];
+      if (!widened) {
+        const size_t numel = source.size() / sizeof(c10::BFloat16);
+        std::vector<uint8_t> converted(numel * sizeof(float));
+        for (size_t i = 0; i < numel; ++i) {
+          store_element(
+              converted.data(),
+              i,
+              static_cast<float>(
+                  load_element<c10::BFloat16>(source.data(), i)));
+        }
+        widened = std::make_shared<OwnedBytes>(
+            OwnedBytes::from_vector(std::move(converted)));
+      }
+      storage = widened;
     }
     const ByteSpan bytes = storage->span();
+    if (bytes.size() != logical_nbytes(device_meta(b.value_id))) {
+      throw std::runtime_error(
+          "vulkan: constant '" + b.key + "' does not match its device size");
+    }
     auto* context = new std::shared_ptr<OwnedBytes>(storage);
     return executorch::runtime::FreeableBuffer(
         bytes.data(), bytes.size(), release_owned_bytes, context);
@@ -2171,7 +2269,7 @@ class VulkanEngineExecutable final : public EngineExecutable {
       }
       requests.push_back(AllocationRequest{
           .value_id = id,
-          .size_bytes = logical_nbytes(value.tensor_meta()),
+          .size_bytes = logical_nbytes(device_meta(id)),
           .memory_kind = memory_kind(storage_at(id)),
       });
     }

backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp

@@ -74,7 +74,8 @@ size_t lower_rms_norms(Graph& graph) {
         shape.values != std::vector<int64_t>{input_meta.sizes.back()} ||
         weight_meta.sizes != shape.values ||
         (input_meta.dtype != ScalarType::Float &&
-         input_meta.dtype != ScalarType::Half) ||
+         input_meta.dtype != ScalarType::Half &&
+         input_meta.dtype != ScalarType::BFloat16) ||
         weight_meta.dtype != input_meta.dtype ||
         output_meta.dtype != input_meta.dtype ||
         output_meta.sizes != input_meta.sizes) {

backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp

@@ -22,14 +22,14 @@
 namespace ptn::vulkan {
 namespace {
 
-Method make_rms_norm() {
+Method make_rms_norm(ScalarType dtype = ScalarType::Float) {
   Method method;
   method.name = "forward";
   Graph& graph = method.graph;
   graph.values = {
-      Value("input", ScalarType::Float, {2, 4}),
-      Value("weight", ScalarType::Float, {4}),
-      Value("output", ScalarType::Float, {2, 4}),
+      Value("input", dtype, {2, 4}),
+      Value("weight", dtype, {4}),
+      Value("output", dtype, {2, 4}),
   };
   graph.value(1).role = ValueRole::Parameter;
   graph.nodes = {
@@ -91,6 +91,19 @@ TEST(LowerRMSNormTest, ReadsUnnamedInputsByPosition) {
   EXPECT_EQ(lower_rms_norms(method.graph), 1);
 }
 
+TEST(LowerRMSNormTest, RewritesBFloat16Norm) {
+  Method method = make_rms_norm(ScalarType::BFloat16);
+
+  EXPECT_EQ(lower_rms_norms(method.graph), 1);
+
+  const Node& norm = method.graph.node(2);
+  EXPECT_EQ(norm.target, "torch.ops.et_vk.rms_norm.default");
+  ASSERT_EQ(norm.inputs.size(), 3);
+  EXPECT_EQ(norm.inputs[0].arg.as_tensor().id, 0);
+  EXPECT_EQ(norm.inputs[1].arg.as_tensor().id, 1);
+  EXPECT_NO_THROW(validate_graph(method.graph));
+}
+
 TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) {
   Method method = make_rms_norm();
   method.graph.node(2).inputs[1].arg = IntListArg{{8}};

backends/native/test/runtime/vulkan/test_vulkan_engine.cpp

@@ -2833,6 +2833,81 @@ std::vector<uint8_t> stored_floats(const std::vector<StoredFloats>& constants) {
   return tensors;
 }
 
+// aten.rms_norm over the last dim of a [2, 4] input with a stored weight.
+std::vector<uint8_t> make_rms_norm_program(fbs::ScalarType dtype) {
+  flatbuffers::FlatBufferBuilder builder;
+  const auto meta = create_tensor_meta(builder, dtype, {2, 4});
+  const auto weight_meta = create_tensor_meta(builder, dtype, {4});
+  const std::vector<flatbuffers::Offset<fbs::TensorValue>> tensor_values = {
+      fbs::CreateTensorValueDirect(builder, "input", meta),
+      fbs::CreateTensorValueDirect(builder, "weight", weight_meta),
+      fbs::CreateTensorValueDirect(builder, "normed", meta),
+  };
+  const std::vector<flatbuffers::Offset<fbs::Output>> input_outputs = {
+      fbs::CreateOutputDirect(builder, "input")};
+  const std::vector<flatbuffers::Offset<fbs::Output>> weight_outputs = {
+      fbs::CreateOutputDirect(builder, "weight")};
+  const std::vector<flatbuffers::Offset<fbs::Output>> norm_outputs = {
+      fbs::CreateOutputDirect(builder, "normed")};
+  const auto eps = fbs::CreateFloatArg(builder, 1e-5);
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> norm_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "input", create_tensor_arg(builder, "input")),
+      fbs::CreateNamedArgumentDirect(
+          builder, "normalized_shape", create_int_list_arg(builder, {4})),
+      fbs::CreateNamedArgumentDirect(
+          builder, "weight", create_tensor_arg(builder, "weight")),
+      fbs::CreateNamedArgumentDirect(
+          builder,
+          "eps",
+          fbs::CreateArgument(
+              builder, fbs::ArgumentValue::FloatArg, eps.Union())),
+  };
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> output_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "", create_tensor_arg(builder, "normed"))};
+  const std::vector<flatbuffers::Offset<fbs::Node>> nodes = {
+      fbs::CreateNodeDirect(
+          builder,
+          "input",
+          fbs::OpKind::PLACEHOLDER,
+          "",
+          nullptr,
+          &input_outputs),
+      fbs::CreateNodeDirect(
+          builder,
+          "weight",
+          fbs::OpKind::PLACEHOLDER,
+          "",
+          nullptr,
+          &weight_outputs),
+      fbs::CreateNodeDirect(
+          builder,
+          "rms_norm",
+          fbs::OpKind::CALL_FUNCTION,
+          "torch.ops.aten.rms_norm.default",
+          &norm_inputs,
+          &norm_outputs),
+      fbs::CreateNodeDirect(
+          builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs),
+  };
+  const auto graph = fbs::CreateGraph(
+      builder,
+      builder.CreateVector(nodes),
+      create_strings(builder, {"input"}),
+      create_strings(builder, {"normed"}),
+      builder.CreateVector(tensor_values));
+  const std::vector<flatbuffers::Offset<fbs::NamedTensorRef>> constants = {
+      fbs::CreateNamedTensorRefDirect(
+          builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER),
+  };
+  return finish_program(
+      builder,
+      graph,
+      {fbs::CreateOutputSpecDirect(builder, "normed")},
+      constants);
+}
+
 constexpr int64_t kHfTokens = 4;
 constexpr int64_t kHfHeads = 4;
 constexpr int64_t kHfKvHeads = 2;
@@ -3439,6 +3514,79 @@ TEST(VulkanEngineTest, PreservesLogicalLongDtypeAndChecksNarrowing) {
       std::runtime_error);
 }
 
+TEST(VulkanEngineTest, RunsBFloat16AddInFloatAndRoundsOutputToNearestEven) {
+  CompiledProgram compiled =
+      compile_program(make_add_program(fbs::ScalarType::BFLOAT16));
+  EXPECT_EQ(compiled.executable->input_dtype(0), ScalarType::BFloat16);
+  EXPECT_EQ(compiled.executable->output_dtype(0), ScalarType::BFloat16);
+
+  // bf16 bits: {1, 2, 3, 1} + {10, 20, 30, 2^-8}.
+  const std::vector<uint16_t> left = {0x3F80, 0x4000, 0x4040, 0x3F80};
+  const std::vector<uint16_t> right = {0x4120, 0x41A0, 0x41F0, 0x3B80};
+  compiled.executable->set_input(
+      0, left.data(), left.size(), ScalarType::BFloat16);
+  compiled.executable->set_input(
+      1, right.data(), right.size(), ScalarType::BFloat16);
+  compiled.executable->execute();
+  std::vector<uint16_t> output(4);
+  compiled.executable->get_output(
+      0, output.data(), output.size(), ScalarType::BFloat16);
+  // {11, 22, 33, 1}: the f32 sum 1 + 2^-8 is a tie and rounds to even.
+  EXPECT_EQ(output, (std::vector<uint16_t>{0x4130, 0x41B0, 0x4204, 0x3F80}));
+}
+
+TEST(VulkanEngineTest, RoundsFloatHostDataOfBFloat16TensorsToBFloat16) {
+  CompiledProgram compiled =
+      compile_program(make_add_program(fbs::ScalarType::BFLOAT16));
+  // 1 + 2^-9 rounds to 1 on input, and the sum 1 + 2^-8 rounds to even on
+  // output.
+  const std::vector<float> left = {1.0f + 0x1p-9f, 1.0f, 3.0f, 1.0f};
+  const std::vector<float> right = {-1.0f, 0x1p-8f, 30.0f, 0.5f};
+  compiled.executable->set_input(
+      0, left.data(), left.size(), ScalarType::Float);
+  compiled.executable->set_input(
+      1, right.data(), right.size(), ScalarType::Float);
+  compiled.executable->execute();
+  std::vector<float> output(4);
+  compiled.executable->get_output(
+      0, output.data(), output.size(), ScalarType::Float);
+  EXPECT_EQ(output, (std::vector<float>{0.0f, 1.0f, 33.0f, 1.5f}));
+  std::vector<int32_t> ints(4);
+  EXPECT_THROW(
+      compiled.executable->get_output(
+          0, ints.data(), ints.size(), ScalarType::Int),
+      std::runtime_error);
+}
+
+TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithStoredWeight) {
+  const std::string header =
+      R"({"weight":{"dtype":"BF16","shape":[4],"data_offsets":[0,8]}})";
+  std::vector<uint8_t> tensors = testing::make_safetensors(header, 8);
+  // bf16 bits: {0.5, 1.5, -2, 3}.
+  const std::vector<uint16_t> weight = {0x3F00, 0x3FC0, 0xC000, 0x4040};
+  std::memcpy(
+      tensors.data() + sizeof(uint64_t) + header.size(),
+      weight.data(),
+      weight.size() * sizeof(uint16_t));
+  CompiledProgram compiled = compile_program(
+      make_rms_norm_program(fbs::ScalarType::BFLOAT16), std::move(tensors));
+  // bf16 bits: rows {2, -2, 2, -2} and {4, 4, -4, 4}, each with an RMS that
+  // normalizes it to +-1.
+  const std::vector<uint16_t> input = {
+      0x4000, 0xC000, 0x4000, 0xC000, 0x4080, 0x4080, 0xC080, 0x4080};
+  compiled.executable->set_input(
+      0, input.data(), input.size(), ScalarType::BFloat16);
+  compiled.executable->execute();
+  std::vector<uint16_t> output(8);
+  compiled.executable->get_output(
+      0, output.data(), output.size(), ScalarType::BFloat16);
+  // {0.5, -1.5, -2, -3} and {0.5, 1.5, 2, 3}.
+  EXPECT_EQ(
+      output,
+      (std::vector<uint16_t>{
+          0x3F00, 0xBFC0, 0xC000, 0xC040, 0x3F00, 0x3FC0, 0x4000, 0x4040}));
+}
+
 TEST(VulkanEngineTest, ExecutesMetadataOnlyAliasAfterAliasLowering) {
   CompiledProgram compiled = compile_program(make_alias_program());
   const std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};