Make TFLite flatbuffer parser backend-neutral.

litert/tensor/internal/BUILD

@@ -20,6 +20,10 @@ package(default_visibility = ["//visibility:public"])
 cc_library(
     name = "mixin",
     hdrs = ["mixin.h"],
+    deps = [
+        ":graph",
+        ":type_id",
+    ],
 )
 
 cc_library(

litert/tensor/internal/mixin.h

@@ -16,19 +16,55 @@ limitations under the License.
 #ifndef LITERT_TENSOR_INTERNAL_MIXIN_H_
 #define LITERT_TENSOR_INTERNAL_MIXIN_H_
 
+#include <memory>
+#include <tuple>
+#include <type_traits>
+
+#include "litert/tensor/internal/graph.h"
+#include "litert/tensor/internal/type_id.h"
+
 namespace litert::tensor {
 
 template <class Mixin>
 class TensorMixin {};
 
 namespace graph {
+
+class MixinRegistrar {
+ public:
+  virtual ~MixinRegistrar() = default;
+  virtual void Register(std::shared_ptr<Operation> op) = 0;
+};
+
 // Provides custom behaviour to operations.
 //
 // - Op is the operation that is being specialized.
 // - Mixin is a tag to identify the mix-in.
 template <class Op, class Mixin>
 class OpMixin {};
 
+template <class MixinTag, class... Ops>
+bool TryRegisterMixinHelper(const std::shared_ptr<Operation>& op,
+                            std::tuple<Ops...>) {
+  bool registered = false;
+  auto try_reg = [&](auto* dummy_op) {
+    using OpType = std::remove_pointer_t<decltype(dummy_op)>;
+    if (op->GetTypeId() == internal::TypeId::Get<OpType>()) {
+      op->extensions.push_back(std::make_unique<OpMixin<OpType, MixinTag>>());
+      registered = true;
+      return true;
+    }
+    return false;
+  };
+  (try_reg(static_cast<Ops*>(nullptr)) || ...);
+  return registered;
+}
+
+template <class MixinTag, class OpsTuple>
+void RegisterMixin(std::shared_ptr<Operation> op) {
+  TryRegisterMixinHelper<MixinTag>(op, OpsTuple{});
+}
+
 }  // namespace graph
 }  // namespace litert::tensor