1#pragma once
2
3#include <unordered_map>
4
5#include "taichi/program/kernel.h"
6#include "taichi/ir/snode.h"
7
8namespace taichi::lang {
9
10class 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 */
18class 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