diff --git a/litert/tensor/internal/BUILD b/litert/tensor/internal/BUILD index f80a54394b7..18e9fe431ec 100644 --- a/litert/tensor/internal/BUILD +++ b/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( diff --git a/litert/tensor/internal/mixin.h b/litert/tensor/internal/mixin.h index d377c030b2c..e28dc722e1b 100644 --- a/litert/tensor/internal/mixin.h +++ b/litert/tensor/internal/mixin.h @@ -16,12 +16,26 @@ limitations under the License. #ifndef LITERT_TENSOR_INTERNAL_MIXIN_H_ #define LITERT_TENSOR_INTERNAL_MIXIN_H_ +#include +#include +#include + +#include "litert/tensor/internal/graph.h" +#include "litert/tensor/internal/type_id.h" + namespace litert::tensor { template class TensorMixin {}; namespace graph { + +class MixinRegistrar { + public: + virtual ~MixinRegistrar() = default; + virtual void Register(std::shared_ptr op) = 0; +}; + // Provides custom behaviour to operations. // // - Op is the operation that is being specialized. @@ -29,6 +43,28 @@ namespace graph { template class OpMixin {}; +template +bool TryRegisterMixinHelper(const std::shared_ptr& op, + std::tuple) { + bool registered = false; + auto try_reg = [&](auto* dummy_op) { + using OpType = std::remove_pointer_t; + if (op->GetTypeId() == internal::TypeId::Get()) { + op->extensions.push_back(std::make_unique>()); + registered = true; + return true; + } + return false; + }; + (try_reg(static_cast(nullptr)) || ...); + return registered; +} + +template +void RegisterMixin(std::shared_ptr op) { + TryRegisterMixinHelper(op, OpsTuple{}); +} + } // namespace graph } // namespace litert::tensor