xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/sdk/program_path.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 //
2 // program_path.h
3 //
4 // Copyright © 2024 Apple Inc. All rights reserved.
5 //
6 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
7 
8 #pragma once
9 
10 #import <string>
11 #import <variant>
12 #import <vector>
13 
14 #import "hash_util.h"
15 
16 namespace executorchcoreml {
17 namespace modelstructure {
18 class Path final {
19 public:
20     static const char* kTypeKeyName;
21     struct Program {
22         /// The type name, used for serializing/deserializing Program type.
23         static const char* kTypeName;
24         struct Function {
25             /// The type name, used for serializing/deserializing Function type.
26             static const char* kTypeName;
27             /// The name key name, used for serializing/deserializing function name.
28             static const char* kNameKeyName;
29             /// Name of the function.
30             std::string name;
31 
FunctionProgram::Function32             Function(std::string name) : name(std::move(name)) { }
33 
34             inline bool operator==(const Function& rhs) const noexcept { return name == rhs.name; }
35 
36             inline bool operator<(const Function& rhs) const noexcept { return name < rhs.name; }
37         };
38 
39         struct Block {
40             /// The type name, used for serializing/deserializing Block type.
41             static const char* kTypeName;
42             /// The index key name, used for serializing/deserializing index.
43             static const char* kIndexKeyName;
44             /// Index of the block.
45             int64_t index;
46 
BlockProgram::Block47             Block(int64_t index) : index(index) { }
48 
49             inline bool operator==(const Block& rhs) const noexcept { return index == rhs.index; }
50 
51             inline bool operator<(const Block& rhs) const noexcept { return index < rhs.index; }
52         };
53 
54         struct Operation {
55             /// The type name, used for serializing/deserializing Operation type.
56             static const char* kTypeName;
57             /// The output key name, used for serializing/deserializing the Operation output.
58             static const char* kOutputKeyName;
59             /// Output name.
60             std::string output_name;
61 
OperationProgram::Operation62             Operation(std::string output_name) : output_name(std::move(output_name)) { }
63 
64             inline bool operator==(const Operation& rhs) const noexcept { return output_name == rhs.output_name; }
65 
66             inline bool operator<(const Operation& rhs) const noexcept { return output_name < rhs.output_name; }
67         };
68 
69         inline bool operator==(const Program& __unused rhs) const noexcept { return true; }
70     };
71 
72     using Component = std::variant<Program, Program::Function, Program::Block, Program::Operation>;
73 
74     /// Appends a component to the path.
75     void append_component(Component component) noexcept;
76 
77     /// Removes the last component.
remove_last_component()78     inline void remove_last_component() noexcept { components_.pop_back(); }
79 
80     /// Removes the first component.
remove_first_component()81     inline void remove_first_component() noexcept { components_.erase(components_.begin()); }
82 
83     /// Returns the number of components.
size()84     inline size_t size() const noexcept { return components_.size(); }
85 
86     /// Returns components.
components()87     inline const std::vector<Component>& components() const noexcept { return components_; }
88 
89     const Component& operator[](size_t index) const { return components_[index]; }
90 
91     inline bool operator==(const Path& rhs) const noexcept { return components() == rhs.components(); }
92 
93 private:
94     std::vector<Component> components_;
95 };
96 }
97 }
98 
99 namespace std {
100 using namespace executorchcoreml::modelstructure;
101 
102 template <> struct hash<Path::Program> {
103     inline size_t operator()(const Path::Program __unused& program) const { return typeid(Path::Program).hash_code(); }
104 };
105 
106 template <> struct hash<Path::Program::Block> {
107     inline size_t operator()(const Path::Program::Block& block) const { return hash<int64_t>()(block.index); }
108 };
109 
110 template <> struct hash<Path::Program::Function> {
111     inline size_t operator()(const Path::Program::Function& function) const {
112         return hash<std::string>()(function.name);
113     }
114 };
115 
116 template <> struct hash<Path::Program::Operation> {
117     inline size_t operator()(const Path::Program::Operation& operation) const {
118         return hash<std::string>()(operation.output_name);
119     }
120 };
121 
122 template <> struct hash<Path> {
123     inline size_t operator()(const Path& path) const { return executorchcoreml::container_hash(path.components()); }
124 };
125 } // namespace std
126