1 | #pragma once |
---|---|
2 | |
3 | #include "taichi/program/callable.h" |
4 | #include "taichi/program/function_key.h" |
5 | |
6 | namespace taichi::lang { |
7 | |
8 | class Program; |
9 | |
10 | class Function : public Callable { |
11 | public: |
12 | enum class IRType { None, AST, InitialIR, OptimizedIR }; |
13 | |
14 | FunctionKey func_key; |
15 | |
16 | Function(Program *program, const FunctionKey &func_key); |
17 | |
18 | // Set the function body to a frontend Python function which generates the C++ |
19 | // AST. |
20 | void set_function_body(const std::function<void()> &func); |
21 | |
22 | // Set the function body to a CHI IR. |
23 | void set_function_body(std::unique_ptr<IRNode> func_body); |
24 | |
25 | [[nodiscard]] std::string get_name() const override; |
26 | |
27 | const std::optional<std::string> &try_get_ast_serialization_data() const { |
28 | return ast_serialization_data_; |
29 | } |
30 | |
31 | void set_ir_type(IRType type) { |
32 | ir_type_ = type; |
33 | } |
34 | |
35 | IRType ir_type() const { |
36 | return ir_type_; |
37 | } |
38 | |
39 | private: |
40 | IRType ir_type_{IRType::None}; |
41 | std::optional<std::string> ast_serialization_data_; // For generating AST-Key |
42 | }; |
43 | |
44 | } // namespace taichi::lang |
45 |