1 | #pragma once |
2 | |
3 | #include <unordered_map> |
4 | |
5 | #include "taichi/program/kernel.h" |
6 | #include "taichi/ir/snode.h" |
7 | |
8 | namespace taichi::lang { |
9 | |
10 | class Program; |
11 | |
12 | /** A mapping from an SNode to its read/write access kernels. |
13 | * |
14 | * The main purpose of this class is to decouple the accessor kernels from the |
15 | * SNode class itself. Ideally, SNode should be nothing more than a group of |
16 | * plain data. |
17 | */ |
18 | class SNodeRwAccessorsBank { |
19 | private: |
20 | struct RwKernels { |
21 | Kernel *reader{nullptr}; |
22 | Kernel *writer{nullptr}; |
23 | }; |
24 | |
25 | public: |
26 | class Accessors { |
27 | public: |
28 | explicit Accessors(const SNode *snode, |
29 | const RwKernels &kernels, |
30 | Program *prog); |
31 | |
32 | // for float and double |
33 | void write_float(const std::vector<int> &I, float64 val); |
34 | float64 read_float(const std::vector<int> &I); |
35 | |
36 | // for int32 and int64 |
37 | void write_int(const std::vector<int> &I, int64 val); |
38 | void write_uint(const std::vector<int> &I, uint64 val); |
39 | int64 read_int(const std::vector<int> &I); |
40 | uint64 read_uint(const std::vector<int> &I); |
41 | |
42 | private: |
43 | const SNode *snode_; |
44 | Program *prog_; |
45 | Kernel *reader_; |
46 | Kernel *writer_; |
47 | }; |
48 | |
49 | explicit SNodeRwAccessorsBank(Program *program) : program_(program) { |
50 | } |
51 | |
52 | Accessors get(SNode *snode); |
53 | |
54 | private: |
55 | Program *const program_; |
56 | std::unordered_map<const SNode *, RwKernels> snode_to_kernels_; |
57 | }; |
58 | |
59 | } // namespace taichi::lang |
60 | |