Add functions for working with dispatch op custom options using the f…

…lex buffer api.

PiperOrigin-RevId: 715918395

tflite/experimental/litert/core/BUILD

@@ -106,6 +106,17 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "dispatch_op_schema",
+    srcs = ["dispatch_op_schema.cc"],
+    hdrs = ["dispatch_op_schema.h"],
+    copts = ["-DFLATBUFFERS_LOCALE_INDEPENDENT=0"],
+    deps = [
+        "//tflite/experimental/litert/cc:litert_buffer_ref",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
 cc_test(
     name = "filesystem_test",
     srcs = ["filesystem_test.cc"],
@@ -138,3 +149,12 @@ cc_test(
 #     ],
 # )
 # copybara:uncomment_end
+
+cc_test(
+    name = "dispatch_op_schema_test",
+    srcs = ["dispatch_op_schema_test.cc"],
+    deps = [
+        ":dispatch_op_schema",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tflite/experimental/litert/core/dispatch_op_schema.cc

@@ -0,0 +1,88 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tflite/experimental/litert/core/dispatch_op_schema.h"
+
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tflite/experimental/litert/cc/litert_buffer_ref.h"
+
+namespace litert {
+namespace internal {
+namespace {
+
+static constexpr const char kBytecodeSizeKey[] = "bytecode_size";
+static constexpr const char kBytecodeOffsetKey[] = "bytecode_offset";
+static constexpr const char kNameKey[] = "name";
+
+}  // namespace
+
+OwningBufferRef<uint8_t> MakeDispatchOpOptions(DispatchOpOptions options) {
+  flexbuffers::Builder fbb;
+
+  // Set maximum width for scalars to 64 bits. This prevents any upsizing of
+  // the buffer when updating the bytecode size and offset in place.
+  fbb.ForceMinimumBitWidth(flexbuffers::BIT_WIDTH_64);
+
+  auto start = fbb.StartMap();
+
+  fbb.UInt(kBytecodeSizeKey, options.bytecode_size);
+  fbb.UInt(kBytecodeOffsetKey, options.bytecode_offset);
+  fbb.String(kNameKey, options.name);
+
+  fbb.EndMap(start);
+  fbb.Finish();
+
+  auto buf = fbb.GetBuffer();
+  OwningBufferRef<uint8_t> res;
+  res.Assign(buf.data(), buf.size());
+
+  return res;
+}
+
+bool UpdateDispatchOpOptionsInPlace(DispatchOpOptions options,
+                                    MutableBufferRef<uint8_t> buffer) {
+  auto opts = flexbuffers::GetRoot(buffer.Data(), buffer.Size()).AsMap();
+
+  // Update name if same len.
+  const auto name_ok = opts[kNameKey].MutateString(options.name);
+
+  // Update bytecode size and offset. Since min scalar bit width is set to max
+  // possible value, it shouldn't fail in theory.
+  const auto size_ok = opts[kBytecodeSizeKey].MutateUInt(options.bytecode_size);
+  const auto offset_ok =
+      opts[kBytecodeOffsetKey].MutateUInt(options.bytecode_offset);
+
+  return name_ok && size_ok && offset_ok;
+}
+
+DispatchOpOptions GetDispatchOpOptions(BufferRef<uint8_t> buffer) {
+  const auto opts = flexbuffers::GetRoot(buffer.Data(), buffer.Size()).AsMap();
+
+  const auto bytecode_size = opts[kBytecodeSizeKey].AsUInt64();
+  const auto bytecode_offset = opts[kBytecodeOffsetKey].AsUInt64();
+  std::string name(opts[kNameKey].AsString().c_str());
+
+  return DispatchOpOptions{
+      bytecode_size,
+      bytecode_offset,
+      std::move(name),
+  };
+}
+
+}  // namespace internal
+}  // namespace litert

tflite/experimental/litert/core/dispatch_op_schema.h

@@ -0,0 +1,59 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_LITERT_CORE_DISPATCH_OP_SCHEMA_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_LITERT_CORE_DISPATCH_OP_SCHEMA_H_
+
+#include "tflite/experimental/litert/cc/litert_buffer_ref.h"
+
+// Utilities for working with the dispatch op custom options buffer. These
+// functions leverage the flexbuffer api under the hood which allows for inplace
+// updates.
+
+namespace litert {
+namespace internal {
+
+// Schema representing the custom options data for dispatch ops. Primarly used
+// to for tracking location of bytecode.
+struct DispatchOpOptions {
+  // The size of the bytecode for the dispatch op.
+  size_t bytecode_size;
+
+  // The offset of the bytecode for the dispatch op relative to the start of the
+  // model file.
+  size_t bytecode_offset;
+
+  // Name of specific dispatch op or entry point to be called in a shared
+  // bytecode module.
+  std::string name;
+};
+
+// Get a serialized representation of the dispatch op options. These should
+// be stored directly in the custom options of the dispatch op.
+OwningBufferRef<uint8_t> MakeDispatchOpOptions(DispatchOpOptions options);
+
+// Update the dispatch op options in the given buffer with the given options.
+// The buffer should be the custom options buffer of the dispatch op. Fails if
+// the passed values would resize the buffer.
+bool UpdateDispatchOpOptionsInPlace(DispatchOpOptions options,
+                                    MutableBufferRef<uint8_t> buffer);
+
+// Get the dispatch op options from the given buffer. The buffer should be the
+// custom options buffer of the dispatch op.
+DispatchOpOptions GetDispatchOpOptions(BufferRef<uint8_t> buffer);
+
+}  // namespace internal
+}  // namespace litert
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_LITERT_CORE_DISPATCH_OP_SCHEMA_H_

tflite/experimental/litert/core/dispatch_op_schema_test.cc

@@ -0,0 +1,74 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tflite/experimental/litert/core/dispatch_op_schema.h"
+
+#include <cstddef>
+
+#include <gtest/gtest.h>
+
+namespace litert {
+namespace internal {
+namespace {
+
+static constexpr size_t kBufferSize = 100;
+static constexpr size_t kBufferOffset = 200;
+static constexpr const char kName[] = "test_name";
+
+TEST(DispatchOpSchemaTest, DispatchOpOptions) {
+  DispatchOpOptions options = {
+      kBufferSize,
+      kBufferOffset,
+      kName,
+  };
+
+  auto buffer = MakeDispatchOpOptions(options);
+  ASSERT_GT(buffer.Size(), 0);
+
+  auto parsed_options = GetDispatchOpOptions(buffer);
+  ASSERT_EQ(parsed_options.bytecode_size, kBufferSize);
+  ASSERT_EQ(parsed_options.bytecode_offset, kBufferOffset);
+  ASSERT_EQ(parsed_options.name, kName);
+}
+
+TEST(DispatchOpSchemaTest, UpdateDispatchOpOptions) {
+  DispatchOpOptions options = {
+      kBufferSize,
+      kBufferOffset,
+      kName,
+  };
+
+  auto buffer = MakeDispatchOpOptions(options);
+  ASSERT_GT(buffer.Size(), 0);
+
+  static constexpr size_t kNewBufferSize = 1000;
+  static constexpr size_t kNewBufferOffset = 2000;
+
+  DispatchOpOptions new_options = {
+      kNewBufferSize,
+      kNewBufferOffset,
+      kName,
+  };
+
+  ASSERT_TRUE(UpdateDispatchOpOptionsInPlace(new_options, buffer));
+
+  auto parsed_options = GetDispatchOpOptions(buffer);
+  ASSERT_EQ(parsed_options.bytecode_size, kNewBufferSize);
+  ASSERT_EQ(parsed_options.bytecode_offset, kNewBufferOffset);
+  ASSERT_EQ(parsed_options.name, kName);
+}
+
+}  // namespace
+}  // namespace internal
+}  // namespace litert